修改模型

import torchvision.models as models
net = models.resnet50()
# 檢視模型定義
print(net)

# output
ResNet(
  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU(inplace=True)
  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  (layer1): Sequential(
    ......
  (fc): Linear(in_features=2048, out_features=1000, bias=True)
)
 

可以看到，Resnet50最後一層(fc)預設輸出1000個節點。
若想將該模型應用於10分類任務中，則需要將最後的輸出節點數修改為10。

import torch.nn as nn
from collections import OrderedDict
# 一層全連線層可能太少，可以再加一層。
classifier = nn.Sequential(OrderedDict([('fc1',nn.Linear(2048,128)),
                          ('relu',nn.ReLU()),
                          ('dropout',nn.Dropout(0.5)),
                          ('fc2',nn.Linear(128,10)),
                          ('output',nn.Softmax(dim=1))]))
# 將net的fc層替換為自定義的classifier
net.fc = classifier
 

再輸出net，可以看到最後一層的fc已修改為定義的內容

  (fc): Sequential(
    (fc1): Linear(in_features=2048, out_features=128, bias=True)
    (relu): ReLU()
    (dropout): Dropout(p=0.5, inplace=False)
    (fc2): Linear(in_features=128, out_features=10, bias=True)
    (output): Softmax(dim=1)
  )

PyTorch模型儲存與讀取

模型儲存

• 模型儲存資料格式：pt
  , pth, pkl

import os
import torch

# 希望使用的GPU編號
os.environ['CUDA_CISIBLE_DEVICES'] = '0'
net.cuda()
# 儲存模型，資料格式可以為 pt, pth, pkl
torch.save(net, './model.pt')
# 儲存權重
torch.save(net.state_dict(), './weight.pt')

模型載入

# 讀取模型
loaded_model = torch.load('./model.pt')
# 將權重載入到模型上，也可先讀取到一個變數中，再為loaded_model賦值，分兩步進行
loaded_model.state_dict = torch.load('./weight.pt')
loaded_model.cuda()

loaded_dict = torch.load('./weight.pt')
print(loaded_dict.keys())
# odict_keys(['conv1.weight', 'bn1.weight', 'bn1.bias', 'bn1.running_mean', 'bn1.running_var', 'bn1.num_batches_tracked', 'layer1.0.conv1.weight', 'layer1.0.bn1.weight', 
......
'fc.fc1.weight', 'fc.fc1.bias', 'fc.fc2.weight', 'fc.fc2.bias'])