从0开始深度学习（21）——读写数据和GPU

学习如何加载和存储权重向量和整个模型，并且将讨论如何利用如何使用单个GPU，然后是如何使用多个GPU和多个服务器（具有多个GPU）进行计算。

1 加载和保存张量

对于单个张量，我们可以使用load和save去加载和保存。

①一个张量

import torch
from torch import nn
from torch.nn import functional as F

x = torch.arange(4)
torch.save(x, 'x-file')
x2 = torch.load('x-file')
x2

②张量列表

y = torch.zeros(4)
torch.save([x, y],'x-files')
x2, y2 = torch.load('x-files')
(x2, y2)

③张量字典

mydict = {'x': x, 'y': y}
torch.save(mydict, 'mydict')
mydict2 = torch.load('mydict')
mydict2

2 加载和保存模型参数

这将保存模型的参数而不是保存整个模型。 例如，如果我们有一个3层多层感知机，我们需要单独指定架构，然后从磁盘加载参数。

以多层感知机为例：

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(20, 256)
        self.output = nn.Linear(256, 10)

    def forward(self, x):
        return self.output(F.relu(self.hidden(x)))

net = MLP()
X = torch.randn(size=(2, 20))
Y = net(X)

我们将模型的参数存储到一个叫"mlp.params"的文件中

torch.save(net.state_dict(), 'mlp.params')

为了恢复模型，我们假设新建了一个同样的模型（这里我直接复制），不需要随机初始化模型参数，而是直接读取文件的模型参数加载进去，由于模型结构一样，模型参数一样，所以在计算相同的X时，结果也应该一样

clone = MLP()
clone.load_state_dict(torch.load('mlp.params'))
clone.eval()

Y_clone = clone(X)
Y_clone == Y

3 GPU

这里我推荐大家使用Kaggle进行示例学习

3.1 查看显卡信息

使用该命令查看信息

!nvidia-smi

可以看到有两块Tesla T4显卡，以及cuda版本号、驱动版本号等等信息

3.2 计算设备

我们可以指定用于存储和计算的设备，如CPU和GPU。 默认情况下，张量是在内存中创建的，然后使用CPU计算它。

在PyTorch中，CPU和GPU可以用torch.device(‘cpu’) 和torch.device(‘cuda’)表示

如果有多个GPU，我们使用torch.device(f’cuda:{i}') 来表示第$i$块GPU（从0开始），可以使用下列代码查看：

torch.device('cpu'), torch.device('cuda'), torch.device('cuda:1')

3.3 查询可用GPU的数量

torch.cuda.device_count()

现在我们定义了两个方便的函数， 这两个函数允许我们在不存在所需所有GPU的情况下运行代码。

def try_gpu(i=0):  #@save
    """如果存在，则返回gpu(i)，否则返回cpu()"""
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

def try_all_gpus():  #@save
    """返回所有可用的GPU，如果没有GPU，则返回[cpu(),]"""
    devices = [torch.device(f'cuda:{i}')
             for i in range(torch.cuda.device_count())]
    return devices if devices else [torch.device('cpu')]

try_gpu(), try_gpu(10), try_all_gpus()

4 张量与GPU

默认情况下，张量是在CPU上创建的

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

4.1 将张量存储在GPU上

我们可以在创建张量时指定存储设备，下面代码是在GPU 0上创建一个张量：

X = torch.ones(2, 3, device=try_gpu())
X

通 3.1 我们知道kaggle提供了2块GPU，我们下面将在第二块GPU上创建张量：

Y = torch.rand(2, 3, device=try_gpu(1))
Y

4.2 复制

我们可以将0号GPU中的数据复制到1号GPU中进行计算：

# 将张量 X 从CPU内存移动到编号为1的GPU设备
Z = X.cuda(1)
print(X)
print(Z)
Y + Z

5 神经网络与GPU

类似地，神经网络模型可以指定设备。 下面的代码将模型参数放在GPU上，当输入为GPU上的张量时，模型将在同一GPU上计算结果：

net = nn.Sequential(nn.Linear(3, 1))
net = net.to(device=try_gpu())
net(X)

让我们确认模型参数存储在同一个GPU上。