pytorch 使用预先训练的模型将图像分类到我选择的50个类别中

vfh0ocws  于 2023-03-02  发布在  其他
关注(0)|答案(1)|浏览(125)

正如我在标题中所说的,我想使用一个预先训练的ML模型,如InceptionV3或类似的模型来对Instagram配置文件中的图像进行分类。
我目前正在使用PyTorch,我可以正确地对我正在尝试的模型进行推理,但我没有找到一个很好的方法来“限制”模型的输出,以匹配我想要的结果,我只有50个类别的选择。
我很想知道如何做这件事。
先谢谢你们了。

8cdiaqws

8cdiaqws1#

您可以通过将预训练版本中的两个层修改为以下内容来实现相同效果:

from torchvision.models import inception_v3

# load the pre-trained model
model = inception_v3(pretrained=True)

# Freeze the weights of the earlier layers since fine-tuning
for param in model.parameters():
    param.requires_grad = False

# Replace the last fully connected layer
inp_feat_lin = model.fc.in_features
model.fc = torch.nn.Linear(inp_feat_lin, 50)

# Replace the AuxLogits layer
inp_feat_aux = model.AuxLogits.fc.in_features
model.AuxLogits.fc = torch.nn.Linear(inp_feat_aux, 50)

# Set the `requires_grad` attribute to `True` for the last layer(s) since fine-tuning
model.fc.requires_grad = True
model.AuxLogits.fc.requires_grad = True

把这个添加到你的类模型的定义中,你可以在pytorch教程页面上深入阅读更多。

相关问题