yixiong/FoodClassifier
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

pretrained=pretrained,解决预训练模型API变化导致的警告。

Browse Source
This commit is contained in:
zhanghuan
2025-09-22 13:59:05 +08:00
parent 13b4e83019
commit 8791b4f2c0
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
net/resnet_embedding.py
+5 -2
View File
@@ -28,8 +28,11 @@ class ResNet50EmbeddingNet(nn.Module):
self.embedding_dim = embedding_dim self.embedding_dim = embedding_dim
self.use_internal_preprocess = use_internal_preprocess self.use_internal_preprocess = use_internal_preprocess
# 加载预训练的ResNet50 # 加载预训练的ResNet50 (使用新的weights参数)
self.backbone = models.resnet50(pretrained=pretrained) if pretrained:
self.backbone = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
else:
self.backbone = models.resnet50(weights=None)
# 获取ResNet50最后一层的输入特征数（2048） # 获取ResNet50最后一层的输入特征数（2048）
backbone_output_dim = self.backbone.fc.in_features backbone_output_dim = self.backbone.fc.in_features