在PyTorch中,你可以使用torch.nn
模块来创建一个神经网络模型。下面是一个简单的例子,展示了如何创建一个包含一个线性层和一个激活函数的简单神经网络:
import torch import torch.nn as nn # 定义一个简单的神经网络模型 class SimpleNN(nn.Module): def __init__(self, input_size, hidden_size, output_size): super(SimpleNN, self).__init__() # 创建一个线性层,输入大小对应于特征数量,输出大小对应于隐藏层大小 self.linear = nn.Linear(input_size, hidden_size) # 创建一个激活函数,这里使用ReLU self.relu = nn.ReLU() # 创建另一个线性层,输入大小对应于隐藏层大小,输出大小对应于输出类别数量 self.output_linear = nn.Linear(hidden_size, output_size) def forward(self, x): # 将输入通过线性层和激活函数 x = self.linear(x) x = self.relu(x) # 将结果通过另一个线性层得到输出 x = self.output_linear(x) return x # 创建一个神经网络实例 input_size = 784 # 假设输入特征数量为784(例如MNIST数据集的图像) hidden_size = 128 output_size = 10 # 假设输出类别数量为10(例如MNIST数据集的10个类别) model = SimpleNN(input_size, hidden_size, output_size) # 打印模型结构 print(model)
在这个例子中,我们定义了一个名为SimpleNN
的类,它继承自nn.Module
。在__init__
方法中,我们创建了两个线性层和一个ReLU激活函数。forward
方法定义了数据在网络中的传播方式,即如何从前一层得到下一层的输出。
你可以根据自己的需求修改这个例子,比如添加更多的层、使用不同的激活函数或者调整网络的结构。