import torch
import torch.nn as nn
import torch.nn.functional as func

class LeNet5(nn.Module):
    def __init__(self,num_classes, grayscale=False):
        """
              num_classes: 分类的数量
              grayscale：是否为灰度图
              """
        super(LeNet5, self).__init__()

        self.grayscale = grayscale
        self.num_classes = num_classes

        if self.grayscale:  # 可以适用单通道和三通道的图像
            in_channels = 1
        else:
            in_channels = 3

        self.conv1 =self.conv1 = nn.Conv2d(in_channels, 6, kernel_size=5)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, num_classes)

    def forward(self, x):
        x = func.max_pool2d(self.conv1(x), 2)   # 原始的模型使用的是 平均池化
        x = func.max_pool2d(self.conv2(x), 2)
        x = x.view(x.size(0), -1)
        x = self.fc3(self.fc2(self.fc1(x)))
        x = func.softmax(x,dim=1)
        return x

#（最后模拟了一个输入，输出一个分类器运算后 10 个 softmax 概率值）
num_classes = 10  # 分类数目
grayscale = True  # 是否为灰度图
data = torch.rand((1, 1, 32, 32))
print("input data:\n", data, "\n")
model = LeNet5(num_classes, grayscale)
x= model(data)
print(x)

2.AlexNet

来自《神经网络与深度学习》

假设输入为32*32大小图像，代码实现与上文所述存在一定差异。

import torch
import torch.nn as nn
class AlexNet(nn.Module):
    def __init__(self,num_classes, grayscale=False):

        super(AlexNet, self).__init__()
        self.grayscale = grayscale
        self.num_classes = num_classes
        if self.grayscale:  # 可以适用单通道和三通道的图像
            in_channels = 1
        else:
            in_channels = 3

        self.features = nn.Sequential(
            nn.Conv2d(in_channels, 96, kernel_size=11,padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2),
            nn.Conv2d(96, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2),
            nn.Conv2d(256, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2),
        )
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(256 * 3 * 3, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, 10),
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), 256 * 3 * 3)
        x = self.classifier(x)
        return x

#最后模拟了一个输入，输出一个分类器运算后的值
num_classes = 10  # 分类数目
grayscale = True  # 是否为灰度图
data = torch.rand((1, 1, 32, 32))
print("input data:\n", data, "\n")
model = AlexNet(num_classes,grayscale)
x=model(data)
print(x)