【PyTorch】torch.save() 和 torch.load() 读写文件

在 PyTorch 中，可以使用 torch 提供的 API 进行文件读写操作，主要用于保存和加载模型、张量数据等。以下是常见的几种文件读写方法。

1. 保存和加载张量

保存张量到文件

可以使用 torch.save() 来将张量保存到文件中。这会将张量保存为一个二进制文件，通常使用 .pt 或 .pth 扩展名。

import torch

# 创建一个张量
x = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])

# 保存张量到文件
torch.save(x, 'tensor.pth')

加载张量

可以使用 torch.load() 来加载之前保存的张量。

# 加载保存的张量
loaded_tensor = torch.load('tensor.pth')

print(loaded_tensor)

输出：

tensor([1.0, 2.0, 3.0, 4.0, 5.0])

2. 保存和加载模型（整个模型）

保存整个模型

在 PyTorch 中，模型的保存有两种方式：

• 保存模型的 状态字典（推荐方式）。
• 保存整个模型（不推荐，除非有特殊需求）。

保存整个模型：

import torch.nn as nn

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 5)

    def forward(self, x):
        return self.fc(x)

# 创建模型实例
model = SimpleModel()

# 保存整个模型
torch.save(model, 'model.pth')

加载整个模型

# 加载整个模型
loaded_model = torch.load('model.pth')

# 使用加载的模型
print(loaded_model)

注意：保存整个模型的方式依赖于保存时的代码，因此这不是推荐的方式。推荐使用 模型的状态字典 来保存模型。

保存模型的状态字典

通常情况下，推荐保存模型的 状态字典（state_dict），这样可以更加灵活地加载模型。

# 保存模型的状态字典
torch.save(model.state_dict(), 'model_state_dict.pth')

加载模型的状态字典

# 创建一个新的模型实例
model2 = SimpleModel()

# 加载状态字典到模型
model2.load_state_dict(torch.load('model_state_dict.pth'))

# 确保模型加载成功
print(model2)

3. 保存和加载优化器的状态字典

在训练过程中，优化器的状态（如动量等）也需要保存和加载，以确保训练的连续性。

保存优化器的状态字典

import torch.optim as optim

# 创建优化器
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 保存优化器的状态字典
torch.save(optimizer.state_dict(), 'optimizer.pth')

加载优化器的状态字典

# 加载优化器的状态字典
optimizer.load_state_dict(torch.load('optimizer.pth'))

4. 使用 torch.utils.data 读写数据

PyTorch 还提供了一个工具库 torch.utils.data 来处理大规模数据的读写，特别是从硬盘读取数据并将其转换为 DataLoader。

使用 TensorDataset 和 DataLoader

from torch.utils.data import TensorDataset, DataLoader

# 创建张量数据
data = torch.randn(100, 3)
targets = torch.randint(0, 2, (100,))

# 创建 TensorDataset
dataset = TensorDataset(data, targets)

# 创建 DataLoader
dataloader = DataLoader(dataset, batch_size=16, shuffle=True)

# 迭代 DataLoader
for batch_data, batch_targets in dataloader:
    print(batch_data, batch_targets)

保存数据到文件

在 PyTorch 中，你可以直接使用 torch.save() 来保存任何数据（如训练集或模型）。

# 保存数据
torch.save(dataset, 'dataset.pth')

加载数据

# 加载数据
loaded_dataset = torch.load('dataset.pth')

# 使用加载的数据
for data, target in DataLoader(loaded_dataset, batch_size=16, shuffle=True):
    print(data, target)

5. 保存和加载多个对象

你可以将多个对象一起保存，例如张量、模型、优化器等，使用 torch.save() 和 torch.load()。

保存多个对象

# 假设有多个对象
x = torch.tensor([1.0, 2.0])
y = torch.tensor([3.0, 4.0])

# 使用字典保存多个对象
torch.save({'x': x, 'y': y}, 'multiple_objects.pth')

加载多个对象

# 加载保存的多个对象
checkpoint = torch.load('multiple_objects.pth')

# 访问加载的对象
x = checkpoint['x']
y = checkpoint['y']

print(x, y)

6. 读取文本文件（非 PyTorch）

如果你需要读取常规的文本文件（例如 CSV 或 TXT 文件），你可以使用 Python 标准库中的 open() 函数或其他库（如 pandas）。

读取文本文件

with open('example.txt', 'r') as file:
    content = file.read()
    print(content)

保存文本文件

with open('output.txt', 'w') as file:
    file.write("Hello, PyTorch!")

总结

• 使用 torch.save() 和 torch.load() 可以方便地保存和加载张量、模型、优化器的状态字典等。
• 对于模型的保存，推荐保存模型的状态字典（state_dict()），而不是整个模型。
• 在处理数据集时，可以利用 TensorDataset 和 DataLoader 配合 torch.save() 和 torch.load() 来保存和加载数据集。
• 对于普通文本文件的读写，可以使用 Python 的内建文件操作方法（如 open()）。