温馨提示×

PyTorch中的卷积神经网络怎么实现

小亿
84
2024-03-05 20:20:02
栏目: 编程语言

在PyTorch中,可以使用torch.nn模块中的Conv2d类来实现卷积神经网络。以下是一个简单的示例,展示如何在PyTorch中实现一个简单的卷积神经网络:

import torch
import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1)
        self.fc1 = nn.Linear(32 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.pool(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = self.pool(x)
        x = x.view(-1, 32 * 7 * 7)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 创建一个SimpleCNN实例
model = SimpleCNN()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 训练模型
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

在上面的示例中,我们定义了一个简单的卷积神经网络模型SimpleCNN,其中包括两个卷积层、ReLU激活函数、最大池化层以及两个全连接层。我们定义了模型的前向传播过程在forward方法中,并在训练过程中使用交叉熵损失函数和Adam优化器来训练模型。

通过这种方式,我们可以使用PyTorch来实现、训练和调整卷积神经网络模型。

0