（3）pokeman_使用resnet18预训练模型

chencaw

已于 2022-10-12 07:47:30 修改

阅读量377

点赞数

分类专栏： torch 文章标签：深度学习 pytorch python

于 2022-10-11 12:31:28 首次发布

本文链接：https://blog.csdn.net/chencaw/article/details/127261183

版权

1、引入模型

from    torchvision.models import resnet18

2、使用预训练模型

（1）pretrained=True为使用预训练

（2）*list展开的意思

    trained_model = resnet18(pretrained=True)
    model = nn.Sequential(*list(trained_model.children())[:-1], #[b, 512, 1, 1]
                          Flatten(), # [b, 512, 1, 1] => [b, 512]
                          nn.Linear(512, 5)
                          ).to(device)

（3）最后一层改为了自己的nn.Linear(512,class_numm)

3、实现代码

import  torch
from    torch import optim, nn
import  visdom
# import  torchvision
from    torch.utils.data import DataLoader
from    torchvision import transforms,datasets
# from    utils import Flatten

# from    pokemon import Pokemon
# from    resnet import ResNet18
from    torchvision.models import resnet18
from    PIL import Image


batchsz = 32
lr = 1e-3
epochs = 10
img_resize = 224


# device = torch.device('cuda')
device = torch.device('cpu')
torch.manual_seed(1234)

# tf = transforms.Compose([
#                      transforms.Resize((224,224)),
#                      transforms.ToTensor(),
#      ])
#输入应该是PIL.Image类型
tf = transforms.Compose([
    #匿名函数
    # lambda x:Image.open(x).convert('RGB'), # string path= > image data
    transforms.Resize(