本文记录利用resnet18预训练模型进行迁移学习,在自己的训练数据集上进行重新训练。相关代码重点部分分别介绍如下:
model=torchvision.models.resnet18(pretrained=True)
num_features=model.fc.in_features
model.fc=nn.Linear(num_features,num_classes)
model=model.to(device)
说明:上述代码首先加载resnet18预训练模型,然后根据训练数据集中的分类数量num_classes修改模型的输出
建议如果有NVIDIA GPU则训练时尽量使用GPU,否则训练时间太久。
如果是使用GPU,则请指定device为GPU,例如:
device=torch.device("cuda" if torch.cuda.is_available() else 'cpu')
然后将model以及输入数据都是用to函数指定到GPU上进行运算:
model=model.to(device)
for i,(images,labels) in enumerate(train_loader):
images=images.to(device)
labels=labels.to(device)
需要指出的是,使用GPU训练出来的模型,如果后续用于分类的时候使用的是CPU,则需要将模型进行转化,否则会报错:
torch.load('tensors.pt', map_location='cpu')
以上为利用resnet18预训练模型训练自己的数据集形成新的特征模型的关键代码,如需完整代码,请关注如下公众号 健哥聊量化,关注之后直接键入: resnet18 即可得到下载链接。该代码中分为两个文件,一个是模型训练的源代码,一个是利用该模型进行分类的源代码。
-------------------- 正文到此结束------------------------
推荐一个公众号:健哥聊量化,会持续推出股票相关基础知识,以及python实现的一些基本的分析代码。欢迎大家关注,二维码如下:
相关文章列表如下: