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的参数