nn.Module模块的使用

特点

每一层都是nn.Module
nn.Module都嵌套在nn.Module中

常见的模块

• Linear
• ReLU
• Sigmoid
• Conv2d
• ConvTransposed2d
• Dropout
• etc.

Container容器

• nn.Sequential()

parameters参数管理

• net.parameters()
• net.named_parameters()

to(device)

可以非常方便地把所有的操作从cpu转移到gpu上

device=torch.device('cuda')
net=Net()
net.to(device)

save and load

由于网络可能训练时间过长，可能中途会断开，这时需要在网络中每隔一段时间就自动保存训练参数，之后再训练时，可以不需要从头开始训练

net.load_state_dict(torch.load('ckpt.mdl')) # load
torch.save(net.state_dict(),'ckpt.mdl')  # save 

train/test

网络模式的切换

net.train() # train
net.eval() # test

定义自己的层

class MyLinear(nn.Module):
	def __init__(self,inp,outp):
		super(MyLinear,self).__init__()
		# requires_grad=True
		self.w=nn.Parameter(torch.randn(outp,inp))
		self.b=nn.Parameter(torch.rand(outp))
	def forward(self,x):
		x=x @ self.w.t()+self.b
		return x