pytorch全连接神经网络如何保存 - 问答

要保存 PyTorch 中的全连接神经网络，您可以使用 torch.save() 函数。这个函数允许您将模型的状态字典保存到本地文件系统中。以下是一个简单的示例，展示了如何保存一个全连接神经网络：

import torch
import torch.nn as nn

# 定义一个简单的全连接神经网络
class FullyConnectedNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(FullyConnectedNN, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

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

# 创建网络实例并设置参数
input_size = 784
hidden_size = 128
output_size = 10
model = FullyConnectedNN(input_size, hidden_size, output_size)

# 定义损失函数和优化器（这里使用随机生成的数据作为示例）
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 训练网络（这里省略了训练过程）

# 保存模型到本地文件
model_save_path = 'fully_connected_nn.pth'
torch.save(model.state_dict(), model_save_path)

在这个示例中，我们首先定义了一个简单的全连接神经网络 FullyConnectedNN，然后创建了一个该网络类型的实例 model。接下来，我们使用随机生成的数据作为示例，定义了损失函数和优化器。最后，我们使用 torch.save() 函数将模型的状态字典保存到本地文件 fully_connected_nn.pth 中。

0 赞

0 踩