PyTorch 多GPU下模型的保存与加载(踩坑笔记)

PyTorch是一个开放源码的机器学习库，支持多GPU并行计算。在使用多GPU训练模型时，保存和加载模型需要特别注意。下面是“PyTorch 多GPU下模型的保存与加载(踩坑笔记)”的攻略过程，具体包含以下几个步骤：

1. 引入必要的库

在保存和加载模型之前，我们需要引入必要的库来支持模型的保存和加载。

import torch
from torch.nn.parallel import DistributedDataParallel as DDP

2. 初始化模型

在使用多GPU训练模型时，通常需要使用DDP包装器对模型进行初始化。DDP是一个用于分布式数据并行处理的包装器，可以在多个GPU上并行计算。

# 初始化模型和DDP包装器
model = MyModel()
model = DDP(model)

3. 保存模型

在保存模型时，需要注意保存DDP包装器的状态，否则在加载模型时可能会导致出错。

# 保存模型
torch.save({
            'model': model.module.state_dict(),
            'optimizer': optimizer.state_dict(),
            'ddp': model.state_dict()
            }, checkpoint_path)

在这里，我们使用torch.save函数保存了模型、优化器和DDP包装器的状态。model.state_dict()返回DDP包装器的状态，而model.module.state_dict()返回模型的状态。这里需要注意，在保存模型时需要使用model.module来获取模型状态，否则会保存失败。

4. 加载模型

在加载模型时需要注意，需要先加载模型和DDP包装器的状态，再将模型和DDP包装器捆绑在一起。

# 加载模型
checkpoint = torch.load(checkpoint_path)
model_state_dict = checkpoint['model']
ddp_state_dict = checkpoint['ddp']
optimizer_state_dict = checkpoint['optimizer']
model = MyModel()
model.load_state_dict(model_state_dict)
model = DDP(model)
model.load_state_dict(ddp_state_dict)
optimizer.load_state_dict(optimizer_state_dict)

在这里，我们使用了torch.load函数加载了模型、优化器和DDP包装器的状态。然后，我们使用model.load_state_dict加载了模型的状态，再将model和DDP包装器捆绑在一起，最后使用optimizer.load_state_dict加载了优化器的状态。

示例1

下面是一个示例，展示了如何在多GPU上训练一个模型，并保存和加载该模型。

import torch
from torch.nn.parallel import DistributedDataParallel as DDP

# 定义模型
class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = torch.nn.Linear(10, 5)
        self.fc2 = torch.nn.Linear(5, 2)
        self.relu = torch.nn.ReLU()

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 初始化模型和DDP包装器
model = MyModel()
model = DDP(model)

# 定义优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 训练模型
for i in range(100):
    x = torch.randn(4, 10)
    y = torch.randint(0, 2, size=(4,))
    optimizer.zero_grad()
    output = model(x)
    loss = torch.nn.functional.cross_entropy(output, y)
    loss.backward()
    optimizer.step()

# 保存模型
checkpoint_path = 'model.pth'
torch.save({
            'model': model.module.state_dict(),
            'optimizer': optimizer.state_dict(),
            'ddp': model.state_dict()
            }, checkpoint_path)

# 加载模型
checkpoint = torch.load(checkpoint_path)
model_state_dict = checkpoint['model']
ddp_state_dict = checkpoint['ddp']
optimizer_state_dict = checkpoint['optimizer']
model = MyModel()
model.load_state_dict(model_state_dict)
model = DDP(model)
model.load_state_dict(ddp_state_dict)
optimizer.load_state_dict(optimizer_state_dict)

示例2

下面是另一个示例，展示了如何在多GPU上训练一个模型，并将模型参数保存为可读文本，以便于人们查看和使用。

import torch
from torch.nn.parallel import DistributedDataParallel as DDP

# 定义模型
class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = torch.nn.Linear(10, 5)
        self.fc2 = torch.nn.Linear(5, 2)
        self.relu = torch.nn.ReLU()

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 初始化模型和DDP包装器
model = MyModel()
model = DDP(model)

# 定义优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 训练模型
for i in range(100):
    x = torch.randn(4, 10)
    y = torch.randint(0, 2, size=(4,))
    optimizer.zero_grad()
    output = model(x)
    loss = torch.nn.functional.cross_entropy(output, y)
    loss.backward()
    optimizer.step()

# 保存模型参数为文本
checkpoint_path = 'model.txt'
with open(checkpoint_path, 'w') as f:
    f.write(str(model.module.state_dict()))

# 加载模型参数
with open(checkpoint_path, 'r') as f:
    state_dict = eval(f.read())
model.load_state_dict(state_dict)

在这个示例中，我们使用了python的文件操作来将模型参数保存为可读文本。加载模型时，我们使用eval函数将文本转换为字典类型，并使用model.load_state_dict函数加载模型参数。