首先,我们需要定义一个神经网络模型。以下是一个示例,展示了如何使用 PyTorch 创建一个简单的卷积神经网络(Convolutional Neural Network,CNN):
Python
import torch import torch.nn as nn import torch.nn.functional as F class Net(nn.Module): def __init__(self): super(Net, self).__init__() # 第一个卷积层:1个输入通道,6个输出通道,5x5的卷积核 self.conv1 = nn.Conv2d(1, 6, 5) # 第二个卷积层:6个输入通道,16个输出通道,5x5的卷积核 self.conv2 = nn.Conv2d(6, 16, 5) # 全连接层:输入维度为 16*5*5,输出维度为 120 self.fc1 = nn.Linear(16 * 5 * 5, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10) def forward(self, input): # Convolutional Layer C1 c1 = F.relu(self.conv1(input)) # Subsampling Layer S2 s2 = F.max_pool2d(c1, (2, 2)) # Convolutional Layer C3 c3 = F.relu(self.conv2(s2)) # Subsampling Layer S4 s4 = F.max_pool2d(c3, 2) # Flatten s4 = torch.flatten(s4, 1) # Fully Connected Layer F5 f5 = F.relu(self.fc1(s4)) # Fully Connected Layer F6 f6 = F.relu(self.fc2(f5)) # Output Layer output = self.fc3(f6) return output # 创建一个实例 net = Net() print(net)
这个神经网络包含了两个卷积层、两个池化层和三个全连接层。你可以根据自己的需求进行修改和扩展。