在PyTorch中搭建卷积神经网络(Convolutional Neural Network,CNN)通常需要使用torch.nn
模块。以下是一个简单的示例,展示了如何使用PyTorch构建一个包含两个卷积层和两个全连接层的CNN:
import torch import torch.nn as nn import torch.nn.functional as F class CNN(nn.Module): def __init__(self): super(CNN, self).__init__() # 第一个卷积层 self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3) # 第二个卷积层 self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3) # 全连接层 self.fc1 = nn.Linear(32 * 5 * 5, 128) self.fc2 = nn.Linear(128, 10) # 假设输出类别数为10 def forward(self, x): x = F.relu(self.conv1(x)) x = F.max_pool2d(x, 2) x = F.relu(self.conv2(x)) x = F.max_pool2d(x, 2) x = x.view(-1, 32 * 5 * 5) x = F.relu(self.fc1(x)) x = self.fc2(x) return x # 实例化模型 model = CNN()
以上代码定义了一个简单的CNN模型,包含两个卷积层和两个全连接层。在forward
方法中,定义了模型的前向传播过程,通过卷积操作、激活函数和池化操作构建了CNN的基本结构。最后实例化了这个CNN模型。
你可以根据自己的需求和问题场景修改模型结构、调整参数,来进一步优化和定制化你的CNN模型。PyTorch提供了丰富的工具和API,可以帮助你更好地构建和训练卷积神经网络。