在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,可以幫助你更好地構建和訓練卷積神經網絡。