深入浅出PyTorch 第五章：PyTorch模型定义

 5.1 PyTorch模型定义的方式

5.1.1 必要的知识回顾

• yTorch模型定义应包括两个主要部分：各个部分的初始化（__init__）；数据流向定义（forward）

• 基于nn.Module，我们可以通过Sequential，ModuleList和ModuleDict三种方式定义PyTorch模型。

5.1.2 Sequential

例子

import torch.nn as nn
net = nn.Sequential(
        nn.Linear(784, 256),
        nn.ReLU(),
        nn.Linear(256, 10), 
        )
print(net)

各层顺序固定，不灵活！

5.1.3 ModuleList

例子

net = nn.ModuleList([nn.Linear(784, 256), nn.ReLU()])
net.append(nn.Linear(256, 10)) # # 类似List的append操作
print(net[-1])  # 类似List的索引访问
print(net)

各层顺序不固定。

5.1.4 ModuleDict

例子

net = nn.ModuleDict({
    'linear': nn.Linear(784, 256),
    'act': nn.ReLU(),
})
net['output'] = nn.Linear(256, 10) # 添加
print(net['linear']) # 访问
print(net.output)
print(net)

各层顺序不固定，且可以赋予层名称。

5.2 利用模型块快速搭建复杂网络

通过构建小的block来搭建复杂的神经网络可以有效地降低代码的复杂度。

5.3 PyTorch修改模型

5.4 PyTorch模型保存与读取

模型存储的两种方法：

# 保存整个模型
torch.save(model, save_dir)
# 保存模型权重
torch.save(model.state_dict, save_dir)

单卡情况的保存与读取：

import os
import torch
from torchvision import models

os.environ['CUDA_VISIBLE_DEVICES'] = '0'   #这里替换成希望使用的GPU编号
model = models.resnet152(pretrained=True)
model.cuda()

# 保存+读取整个模型
torch.save(model, save_dir)
loaded_model = torch.load(save_dir)
loaded_model.cuda()

# 保存+读取模型权重
torch.save(model.state_dict(), save_dir)
loaded_dict = torch.load(save_dir)
loaded_model = models.resnet152()   #注意这里需要对模型结构有定义
loaded_model.state_dict = loaded_dict
loaded_model.cuda()