要保存 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
中。