diff --git a/net/resnet_embedding.py b/net/resnet_embedding.py index 9f8175d..7d8c952 100644 --- a/net/resnet_embedding.py +++ b/net/resnet_embedding.py @@ -28,8 +28,11 @@ class ResNet50EmbeddingNet(nn.Module): self.embedding_dim = embedding_dim self.use_internal_preprocess = use_internal_preprocess - # 加载预训练的ResNet50 - self.backbone = models.resnet50(pretrained=pretrained) + # 加载预训练的ResNet50 (使用新的weights参数) + if pretrained: + self.backbone = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1) + else: + self.backbone = models.resnet50(weights=None) # 获取ResNet50最后一层的输入特征数(2048) backbone_output_dim = self.backbone.fc.in_features