本篇主要介绍所有网络层次类的父类:nn.module

一、内置一些网络层

  • Linear
  • ReLu
  • Sigmoid
  • Conv2d
  • ConvTransposed2d
  • Dropout

二、Container

  即Sequential

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)

三、parameters

图片描述

四、modules

  • modules: 指所有的节点
  • children:指直系孩子
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,然后他只有一个childrensequential,他的modules则如下面的输出所示:需要注意的是,自己本身是0号module

    for name, t in net.named_parameters():
        print('parameters:', name, t.shape)

    for name, m in net.named_children():
        print('children:', name, m)
        
    for name, m in net.named_modules():
        print('modules:', name, m)
'''
parameters: net.0.net.weight torch.Size([3, 4])
parameters: net.0.net.bias torch.Size([3])
parameters: net.2.weight torch.Size([2, 3])
parameters: net.2.bias torch.Size([2])
children: net Sequential(
  (0): BasicNet(
    (net): Linear(in_features=4, out_features=3, bias=True)
  )
  (1): ReLU()
  (2): Linear(in_features=3, out_features=2, bias=True)
)
modules:  Net(
  (net): Sequential(
    (0): BasicNet(
      (net): Linear(in_features=4, out_features=3, bias=True)
    )
    (1): ReLU()
    (2): Linear(in_features=3, out_features=2, bias=True)
  )
)
modules: net Sequential(
  (0): BasicNet(
    (net): Linear(in_features=4, out_features=3, bias=True)
  )
  (1): ReLU()
  (2): Linear(in_features=3, out_features=2, bias=True)
)
modules: net.0 BasicNet(
  (net): Linear(in_features=4, out_features=3, bias=True)
)
modules: net.0.net Linear(in_features=4, out_features=3, bias=True)
modules: net.1 ReLU()
modules: net.2 Linear(in_features=3, out_features=2, bias=True)
'''

五、to(device)

  可以非常方便地把网络上的运算转移到指定设备(GPU/CPU

device = torch.device('cuda')
net = Net()
net.to(device)
# 需要注意的是,对于nn.module(比如这里的net)来说
# net = net.to(device)和net.to(device)是等价的
# 但对于tensor,a = a.to(device)和a.to(device)是不同的,返回的是不同的东西

六、save and load

  可以保存和加载中间状态,便于长时间可能间断的训练

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

七、train/test

  切换train/test的状态,前面的文章也提到过,在训练的时候是需要更新梯度等信息而test的时候是不需要的,因此需要作状态的切换

net.train()
net.eval()

八、实现自己的网络

  实现自己类的很大程度上的意义是某些类比如打平,只有函数而不是类,而sequential里面只能是一些列

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)

1. 一个自己的线性层的例子

  要知道的是,nn.module里是有线性层的,这里只是举例说明怎么创建一个自己的类

class MyLinear(nn.Module):

    def __init__(self, inp, outp):
        super(MyLinear, self).__init__()

        # requires_grad = True
        # 经过nn.Parameter打包之后的tensor就可以在之后被优化更新
        # 否则就需要加上上面的需要梯度信息的声明。
        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