pytorch基础知识整理（四） 模型

1，模型构造模板

torch.nn.Module()是所有网络模型的基类，所有网络都需要继承此类，模板如下：

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__() #'表示继承父类的__init__()方法'
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
       x = F.relu(self.conv1(x))
       return F.relu(self.conv2(x))

2，torch.nn.Module()的常用方法

model.modules() ， model.children() 返回模型所有模块/子模块的迭代器。
model.named_modules()， model.named_children() 返回模型所有模块/子模块的名字和模块本身的迭代器。
model.parameters() 返回模型所有参数的迭代器。常用来作为optimizer的迭代器。
model.register_parameter(name, param) 向模型添加parameter。
model.register_buffer(name, tensor) 向模型添加buffer。
model.state_dict(), model.load_state_dict() 返回/加载 状态字典。
model.train(), model.eval() 训练/推理模式，仅影响模型中的dropout和BN层。
model.cpu(), model.cuda() 把模型中的所有parameters和buffers赋值到CPU/GPU中。
model.float(), model.half(), model.double() 转换模型的所有parameters和buffers的类型。
model.zero_grad() 把模型所有parameters的梯度置0，和optimizer.zero_grad()完全等效。
注：可以用model.layer_name的方法得到子模型，如model.conv1.parameters()就得到了conv1的参数