特点
每一层都是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