PyTorch 中使用 __init__ 方法构建 AlexNet 网络

本文将介绍如何使用 PyTorch 中的 __init__ 方法构建经典的 AlexNet 网络。

构建 AlexNet 网络

下面是一个使用 __init__ 方法定义 AlexNet 网络各层的示例代码：

import torch
import torch.nn as nn

class AlexNet(nn.Module):
    def __init__(self, num_classes=1000):
        super(AlexNet, self).__init__()
        
        # 定义卷积层
        self.conv1 = nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2)
        self.conv2 = nn.Conv2d(64, 192, kernel_size=5, padding=2)
        self.conv3 = nn.Conv2d(192, 384, kernel_size=3, padding=1)
        self.conv4 = nn.Conv2d(384, 256, kernel_size=3, padding=1)
        self.conv5 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        
        # 定义池化层
        self.pool = nn.MaxPool2d(kernel_size=3, stride=2)
        
        # 定义全连接层
        self.fc1 = nn.Linear(256 * 6 * 6, 4096)
        self.fc2 = nn.Linear(4096, 4096)
        self.fc3 = nn.Linear(4096, num_classes)
        
    def forward(self, x):
        # 在 forward 方法中定义前向传播的过程
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = torch.relu(self.conv3(x))
        x = torch.relu(self.conv4(x))
        x = self.pool(torch.relu(self.conv5(x)))
        
        x = x.view(x.size(0), -1)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        
        return x

在这个示例中，我们首先继承了 nn.Module 类，并在 __init__ 方法中调用了父类的构造方法。然后，我们定义了模型的各个层，包括 5 个卷积层、3 个池化层和 3 个全连接层。在 forward 方法中，我们定义了前向传播的过程，即通过各个层的组合来计算输出。

这里使用了 PyTorch 提供的常用的层类，如 nn.Conv2d 表示 2D 卷积层，nn.MaxPool2d 表示 2D 最大池化层，nn.Linear 表示全连接层。利用这些层类，我们可以方便地定义模型的各个层。

总结

本文介绍了使用 PyTorch 中的 __init__ 方法构建 AlexNet 网络的步骤。通过这个示例，我们可以了解如何使用 PyTorch 定义模型的各个层，并使用 forward 方法定义模型的前向传播过程。