nn.Module
是在pytorch使用非常广泛的类,搭建网络基本都需要用到这个。
当我们搭建自己的网络时,可以继承官方写好的nn.Module
模块,为什么要用这个呢?好处如下:
Linear、ReLU、Sigmoid、Conv2d、Dropout
不用自己一个一个的写这些函数了,这也是为什么我们用框架的原因之一吧。
比如我们经常用到的 nn.Sequential()
,顾名思义,将网络模块封装在一个容器中,可以方面网络搭建
如下面一个例子:
class TestNet(nn.Module):
def __init__(self):
super(TestNet, self).__init__()
self.net = nn.Sequential(nn.Conv2d(1, 16, stride=1, padding=1),
nn.MaxPool2d(2, 2),
Flatten(),
nn.Linear(1*14*14, 10))
def forward(self, x):
return self.net(x)
参数名字可以自动生成(想想如果自己去命名,百万参数的网络没法搭建),然后这些参数都可以传到优化器里面去优化
class BasicNet(nn.Module):
def __init__(self):
super(BasicNet, self).__init__()
self.net = nn.Linear(4, 3)
def forward(self, x):
return self.net(x)
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.net = nn.Sequential(BasicNet(),
nn.ReLU(),
nn.Linear(3, 2))
def forward(self, x):
return self.net(x)
比如上面的代码,我们可以看出Net网络中有5个孩子节点:nn.Sequential,BasicNet, nn.ReLU,nn.Linear,BasicNet里面的nn.Linear
nn.Module
还有一个功能是将某个网络所有成员、函数、操作都搬移到GPU上面。
采用代码如下:
device = torch.device('cuda')
net = Net()
net.to(device)
上面device代表当前的设备是GPU还是CPU,需要注意的是为什么我们不写
net = net.to(device)
其实效果是一样的,采用nn.Module
模块,net加上.to(device)
,还是net。如果是变量则不是一样的,即如果对于tensor bias
,那么bias
和bias.to(device)
不是一样的,则需要重新命名。
可以方面我们保存和加载模型
加载模型:
net.load_state_dict(torch.load('ckpt.mdl'))
保存模型:
torch.save(net.state_dict(), 'ckpt.mdl')
方便训练和测试进行切换,为什么?因为网络中Dropout和BN在训练和测试是不一样的,需要切换
如果不切换效果就会很差,这个是容易犯的一个错误。
net.train()
net.eval()
官方给的模块还是基础操作的,如果自己要搭建复杂的操作也容易实现,一个典型的例子就是可以自己设计一个新的损失函数。
下面给出将tensor压平的例子(nn.Module
没有这个操作):
class Flatten(nn.Module):
def __init__(self):
super(Flatten, self).__init__()
def forward(self, input):
return input.view(input.size(0), -1)
class TestNet(nn.Module):
def __init__(self):
super(TestNet, self).__init__()
self.net = nn.Sequential(nn.Conv2d(1, 16, stride=1, padding=1),
nn.MaxPool2d(2, 2),
Flatten(), #自己定义的
nn.Linear(1*14*14, 10))
def forward(self, x):
return self.net(x)
Flatten
压平的操作则是我们自己构建的类,可以方便后续BasicNet类使用,注意nn.Sequential
里面必须是类。
且在上面例子中Flatten
不需要接任何参数。
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.randn(outp))
def forward(self, x):
x = x @ self.w.t() + self.b
return x
在上面自己写的线性层 y = w x + b y=wx+b y=wx+b,可以看出 w w w和 b b b必须要使用nn.Parameter
这个模块。原因是只用加上了nn.Parameter
后, w w w和 b b b才可以用优化器SGD等进行优化。
如果不写nn.Parameter
那么则需要写requires_grad = True
,还要自己写优化器,就很麻烦。用了Parameter
可以方便我们优化网络:
model = MyLinear.to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# backprop
optimizer.zero_grad()
loss.backward()
optimizer.step()