Pytorch入门教程学习笔记（四）自定义模型、模型读取和存储、GPU计算

最新推荐文章于 2024-08-02 12:25:21 发布

我就是黑凤梨

最新推荐文章于 2024-08-02 12:25:21 发布

阅读量552

点赞数 2

文章标签： python 深度学习人工智能 pytorch

本文链接：https://blog.csdn.net/hzy199772/article/details/116095572

版权

4.4 自定义Layer

本节将介绍如何使用Module来自定义层，从而可以被重复调用。

4.4.1 不含参数的自定义

自定义的CenteredLayer类通过继承Module类自定义了一个将输入减掉均值后输出的层，并将层的计算定义在了forward函数里。这个层里不含模型参数。

class CenteredLayer(nn.Module):
    def __init__(self, **kwargs):
        super(CenteredLayer, self).__init__(**kwargs)
    def forward(self, x):
        return x- x.mean()

#实例化这个层，然后做前向计算。
layer = CenteredLayer()
layer(torch.tensor([1, 2, 3, 4, 5], dtype=torch.float))

tensor([-2., -1.,  0.,  1.,  2.])

#用它来构造更复杂的模型
net = nn.Sequential(nn.Linear(8, 128), CenteredLayer())
#打印自定义层各个输出的均值。因为均值是浮点数，所以它的值是一个很接近0的数。
y = net(torch.rand(4, 8))
y.mean().item()

-1.862645149230957e-09

4.4.2 含参数的自定义层

我们还可以自定义含模型参数的自定义层。其中的模型参数可以通过训练学出。

在4.2节（模型参数的访问、初始化和共享）中介绍了Parameter类其实是Tensor的子类，如果一个Tensor是Parameter，那么它会自动被添加到模型的参数列表里。所以在自定义含模型参数的层时，我们应该将参数定义成Parameter，除了像4.2.1节那样直接定义成Parameter类外，还可以使用ParameterList和ParameterDict分别定义参数的列表和字典。

ParameterList接收一个Parameter实例的列表作为输入然后得到一个参数列表，使用的时候可以用索引来访问某个参数，另外也可以使用append和extend在列表后面新增参数。

class MyListDense(nn.Module):
    def __init__(self):
        super(MyListDense, self).__init__()
        self.params = nn.ParameterList([nn.Parameter(torch.randn(4,4)) for i in range(3)])
        self.params.append(nn.Parameter(torch.randn(4, 1)))
    def forward(self, x):
        for i in range(len(self.params)):
            x = torch.mm(x, self.params[i])
        return x
net = MyListDense()
print(net)

MyListDense(
  (params): ParameterList(
      (0): Parameter containing: [torch.FloatTensor of size 4x4]
      (1): Parameter containing: [torch.FloatTensor of size 4x4]
      (2): Parameter containing: [torch.FloatTensor of size 4x4]
      (3): Parameter containing: [torch.FloatTensor of size 4x1]
  )
)

ParameterDict接收一个Parameter实例的字典作为输入然后得到一个参数字典，然后可以按照字典的规则使用了。例如使用update()新增参数，使用keys()返回所有键值，使用items()返回所有键值对等等

class MyDictDense(nn.Module):
    def __init__(self):
        super(MyDictDense, self).__init__()
        self.params = nn.ParameterDict({
                'linear1': nn.Parameter(torch.randn(4, 4)),
                'linear2': nn.Parameter(torch.randn(4, 1))
        })
        self.params.update({'linear3': nn.Parameter(torch.randn(4, 2))}) # 新增

    def forward(self, x, choice='linear1'):
        return torch.mm(x, self.params[choice])

net = MyDictDense()
print(net)

MyDictDense(
  (params): ParameterDict(
      (linear1): Parameter containing: [torch.FloatTensor of size 4x4]
      (linear2): Parameter containing: [torch.FloatTensor of size 4x1]
      (linear3): Parameter containing: [torch.FloatTensor of size 4x2]
  )
)

#可以根据传入的键值来进行不同的前向传播
x = torch.ones(1, 4)
print(net(x, 'linear1'))
print(net(x, 'linear2'))
print(net(x, 'linear3'))

tensor([[-1.2218, -1.2171,  1.0751, -2.5647]], grad_fn=<MmBackward>)
tensor([[-1.0242]], grad_fn=<MmBackward>)
tensor([[ 3.3657, -2.4289]], grad_fn=<MmBackward>)

#也可以使用自定义层构造模型。它和PyTorch的其他层在使用上很类似。
net = nn.Sequential(
    MyDictDense(),
    MyListDense(),
)
print(net)
print(net(x))

Sequential(
  (0): MyDictDense(
    (params): ParameterDict(
        (linear1): Parameter containing: [torch.FloatTensor of size 4x4]
        (linear2): Parameter containing: [torch.FloatTensor of size 4x1]
        (linear3): Parameter containing: [torch.FloatTensor of size 4x2]
    )
  )
  (1): MyListDense(
    (params): ParameterList(
        (0): Parameter containing: [torch.FloatTensor of size 4x4]
        (1): Parameter containing: [torch.FloatTensor of size 4x4]
        (2): Parameter containing: [torch.FloatTensor of size 4x4]
        (3): Parameter containing: [torch.FloatTensor of size 4x1]
    )
  )
)
tensor([[-11.1288]], grad_fn=<MmBackward>)

4.5读取和存储

在实际中，我们有时需要把训练好的模型部署到很多不同的设备。在这种情况下，我们可以把内存中训练好的模型参数存储在硬盘上供后续读取使用。

4.5.1 读写Tensor

可以直接使用save函数和load函数分别存储和读取Tensor。save使用Python的pickle实用程序将对象进行序列化，然后将序列化的对象保存到disk，使用save可以保存各种对象,包括模型、张量和字典等。而load使用pickle unpickle工具将pickle的对象文件反序列化为内存。

下面的例子创建了Tensor变量x，并将其存在文件名同为x.pt的文件里。

x = torch.ones(3)
torch.save(x, 'x.pt')

#将数据从存储的文件读回内存。
x2 = torch.load('x.pt')
x2

tensor([1., 1., 1.])

#存储一个Tensor列表并读回内存
y = torch.zeros(4)
torch.save([x, y], 'xy.pt')
xy_list = torch.load('xy.pt')
xy_list

[tensor([1., 1., 1.]), tensor([0., 0., 0., 0.])]

#存储并读取一个从字符串映射到Tensor的字典。
torch.save({'x': x, 'y': y}, 'xy_dict.pt')
xy = torch.load('xy_dict.pt')
xy

{'x': tensor([1., 1., 1.]), 'y': tensor([0., 0., 0., 0.])}

4.5.2 读写Model

4.5.2.1 state_dict

在PyTorch中，Module的可学习参数(即权重和偏差)，模块模型包含在参数中(通过model.parameters()访问)。state_dict是一个从参数名称映射到参数Tesnor的字典对象。

class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.hidden = nn.Linear(3, 2)
        self.act = nn.ReLU()
        self.output = nn.Linear(2, 1)

    def forward(self, x):
        a = self.act(self.hidden(x))
        return self.output(a)

net = MLP()
net.state_dict()

OrderedDict([('hidden.weight',
              tensor([[-0.0796, -0.5565, -0.5132],
                      [ 0.4249, -0.5447,  0.2015]])),
             ('hidden.bias', tensor([ 0.3258, -0.5268])),
             ('output.weight', tensor([[0.1478, 0.3633]])),
             ('output.bias', tensor([-0.3196]))])

只有具有可学习参数的层(卷积层、线性层等)才有state_dict中的条目。优化器(optim)也有一个state_dict，其中包含关于优化器状态以及所使用的超参数的信息。

optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
optimizer.state_dict()

{'state': {},
 'param_groups': [{'lr': 0.001,
   'momentum': 0.9,
   'dampening': 0,
   'weight_decay': 0,
   'nesterov': False,
   'params': [0, 1, 2, 3]}]}

4.5.2.2 保存和加载模型

PyTorch中保存和加载训练模型有两种常见的方法:

1.仅保存和加载模型参数(state_dict)；

2.保存和加载整个模型。

☆1. 保存和加载state_dict(推荐方式)

#保存
torch.save(model.state_dict(), PATH) # 推荐的文件后缀名是pt或pth
#加载
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))

☆2.保存和加载整个模型

#保存
torch.save(model, PATH)
#加载
model = torch.load(PATH)

#采用推荐的方法一来实验一下:
X = torch.randn(2, 3)
Y = net(X)

PATH = "./net.pt"
torch.save(net.state_dict(), PATH)

net2 = MLP()
net2.load_state_dict(torch.load(PATH))
Y2 = net2(X)
Y2 == Y

tensor([[True],
        [True]])

4.6 GPU计算

4.6.1 计算设备

PyTorch可以指定用来存储和计算的设备，如使用内存的CPU或者使用显存的GPU。默认情况下，PyTorch会将数据创建在内存，然后利用CPU来计算。

#用torch.cuda.is_available()查看GPU是否可用:
torch.cuda.is_available() # 输出 True
#查看GPU数量：
torch.cuda.device_count() # 输出 1
#查看当前GPU索引号，索引号从0开始：
torch.cuda.current_device() # 输出 0
#根据索引号查看GPU名字:
torch.cuda.get_device_name(0) # 输出 'GeForce GTX 1650'

'GeForce GTX 1650'

4.6.2 Tensor的GPU计算

默认情况下，Tensor会被存在内存上。因此，之前我们每次打印Tensor的时候看不到GPU相关标识。

使用.cuda()可以将CPU上的Tensor转换（复制）到GPU上。如果有多块GPU，我们用.cuda(i)来表示第 i块GPU及相应的显存（i从0开始）且cuda(0)和cuda()等价。

x = torch.tensor([1, 2, 3])
x

tensor([1, 2, 3])

x = x.cuda(0)
x

tensor([1, 2, 3], device='cuda:0')

#可以通过Tensor的device属性来查看该Tensor所在的设备。
x.device

device(type='cuda', index=0)

#可以直接在创建的时候就指定设备。
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

x = torch.tensor([1, 2, 3], device=device)
# or
x = torch.tensor([1, 2, 3]).to(device)
x

tensor([1, 2, 3], device='cuda:0')

#如果对在GPU上的数据进行运算，那么结果还是存放在GPU上
y = x**2
y

tensor([1, 4, 9], device='cuda:0')

需要注意的是，存储在不同位置中的数据是不可以直接进行计算的。即存放在CPU上的数据不可以直接与存放在GPU上的数据进行运算，位于不同GPU上的数据也是不能直接进行计算的。

4.6.3 Model的GPU计算

同Tensor类似，PyTorch模型也可以通过.cuda转换到GPU上。我们可以通过检查模型的参数的device属性来查看存放模型的设备。

net = nn.Linear(3, 1)
list(net.parameters())[0].device

device(type='cpu')

#其转换到GPU上:
net.cuda()
list(net.parameters())[0].device

device(type='cuda', index=0)

说明

说明：本博客是对如何使用pytorch用于深度学习学习过程的记录和总结。
学习教程为：《动手学深度学习》和https://tangshusen.me/Dive-into-DL-PyTorch/#/
这里推荐这个网址，将动手学深度学习改为了Pytorch实现，很有意义！
代码是借鉴了学习教程并从自己写的Jupyter中导出的，复制进Jupyter可以运行

我就是黑凤梨

关注

2
点赞
踩
0

收藏

觉得还不错? 一键收藏
1
评论
Pytorch入门教程学习笔记（四）自定义模型、模型读取和存储、GPU计算

目录4.4 自定义Layer4.4.1 不含参数的自定义4.4.2 含参数的自定义层4.5读取和存储4.5.1 读写Tensor4.5.2 读写Model4.5.2.1 state_dict4.5.2.2 保存和加载模型4.6 GPU计算4.6.1 计算设备4.6.2 Tensor的GPU计算4.6.3 Model的GPU计算说明4.4 自定义Layer本节将介绍如何使用Module来自定义层，从而可以被重复调用。4.4.1 不含参数的自定义自定义的CenteredLayer类通过继承Module类
复制链接

扫一扫