pytorch是包含一些常见的神经网络模型的,ResNet34、ResNet18、VGG等等,都在models
模块中,调用接口如下:
model = models.resnet34pretrained=True, progress=True)
预训练pretrained
和进度条progress
设为True或False就看你自己需求咯~
⚠️与此同时,你可能还会遇到pytorch官网的输出类别数与你所需不同:你想分七类,可resnet34官网是1000类,那就调整一下呗:
# 修改最后线性层的输出通道数
model.fc = nn.Linearmodel.fc.in_features, 9) # 想输出为9个类别时