您的位置:宽带测速网 > 编程知识 > 如何在PyTorch中创建一个神经网络模型

如何在PyTorch中创建一个神经网络模型

2025-06-23 19:27来源:互联网 [ ]

在PyTorch中创建神经网络模型通常需要定义一个继承自torch.nn.Module类的自定义类。下面是一个简单的示例:

import torchimport torch.nn as nnclass SimpleNN(nn.Module):def __init__(self):super(SimpleNN, self).__init__()self.fc1 = nn.Linear(784, 128)# 定义一个全连接层self.relu = nn.ReLU()# 定义一个激活函数self.fc2 = nn.Linear(128, 10)# 定义另一个全连接层def forward(self, x):x = self.fc1(x)x = self.relu(x)x = self.fc2(x)return x

在上面的示例中,我们定义了一个简单的神经网络模型SimpleNN,包括两个全连接层和一个ReLU激活函数。__init__方法用于定义模型的结构,forward方法用于定义模型的前向传播过程。

要使用这个模型,可以实例化一个对象并传入输入数据进行前向传播计算:

model = SimpleNN()input_data = torch.randn(1, 784)# 创建一个输入数据张量output = model(input_data)# 进行前向传播print(output)

这样就可以在PyTorch中创建一个简单的神经网络模型了。您可以根据自己的需求定义更复杂的模型结构和前向传播过程。