首页 > 其他 > 详细

Pytorch下微调网络模型进行图像分类

时间:2020-05-15 00:56:31      阅读:127      评论:0      收藏:0      [点我收藏+]

利用ImageNet下的预训练权重采用迁移学习策略,能够实现模型快速训练,提高图像分类性能。下面以vgg和resnet网络模型为例,微调最后的分类层进行分类。

 

(1)微调VGG模型进行图像分类(以vgg16为例)

import torch 
import torchvision.models as models

classes_num = 200 # 数据集的类别数

model = models.vgg16(pretrained=True)
model.classifier[-1].out_features = classes_num
model = model.cuda()
print(model)

 

(2)微调ResNet模型进行图像分类(以ResNet-34为例)

import torch
import torchvision.models as models

classes_num = 200 #数据集的类别数

model = models.resnet34(pretrained=True)
fc_features = model.fc.in_features
model.fc = torch.nn.Linear(fc_features, classes_num)
model = model.cuda()
print(model)

  

Pytorch下微调网络模型进行图像分类

原文:https://www.cnblogs.com/rs-xiaosheng/p/12892362.html

(0)
(0)
   
举报
评论 一句话评论(0
关于我们 - 联系我们 - 留言反馈 - 联系我们:wmxa8@hotmail.com
© 2014 bubuko.com 版权所有
打开技术之扣,分享程序人生!