一.AlexNet

1.模型设计

激活函数使用ReLU，计算简单，易于训练

2.pytorch实现AlexNet模型

import torch
import torchvision 
import torchvision.transforms as transforms
from torch import nn

# 超参初始化
lr = 0.001
num_epochs = 5
batch_size = 128

#检查GPU是否可用
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

# 获取MNIST数据集
def load_data_fashion_mnist(batch_size, resize=None):
    # 定义转换，将图像转换为Tensor，并且归一化
    transform = transforms.ToTensor()
    if resize:
        transform = transforms.Compose([
            transforms.Resize(resize),
            transforms.ToTensor()
        ])
    # 加载数据集
    mnist_train = torchvision.datasets.FashionMNIST(root='./database', train=True, download=True, transform=transform)
    mnist_test = torchvision.datasets.FashionMNIST(root='./database', train=False, download=True, transform=transform)

    # 加载数据，windows下只支持num_workers=0，其他可以为4
    train_iter = torch.utils.data.DataLoader(mnist_train, batch_size=batch_size, shuffle=True, num_workers=0)
    test_iter = torch.utils.data.DataLoader(mnist_test, batch_size=batch_size, shuffle=False, num_workers=0)
    return train_iter, test_iter

train_iter, test_iter = load_data_fashion_mnist(batch_size, 224)

# 定义模型
net = nn.Sequential(
    nn.Conv2d(1, 96, kernel_size=11, stride=4, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=3, stride=2),
    nn.Conv2d(96, 256, kernel_size=5, padding=2),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=3, stride=2),
    nn.Conv2d(256, 384, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.Conv2d(384, 384, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.Conv2d(384, 256, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=3, stride=2),
    nn.Flatten(),
    nn.Linear(256 * 5 * 5, 4096),
    nn.ReLU(),
    nn.Dropout(p=0.5),
    nn.Linear(4096, 4096),
    nn.ReLU(),
    nn.Dropout(p=0.5),
    nn.Linear(4096, 10)
)

net.to(device)

# 损失函数
loss = nn.CrossEntropyLoss()

# 优化器
optimizer = torch.optim.Adam(net.parameters(), lr=lr)

# 计算准确率
def evaluate_accuracy(data_iter, net):
    acc_sum, n = 0.0, 0
    with torch.no_grad(): # 临时关闭梯度计算
        for X, y in data_iter:
            X, y = X.to(device), y.to(device)
            net.eval() # 评估模式, 关闭dropout
            acc_sum += (net(X).argmax(dim=1) == y).float().sum().item()
            net.train() # 训练模式
            n += y.shape[0]
    return acc_sum / n

# 定义训练模型函数
def train_net(net, train_iter, test_iter, loss, num_epochs, batch_size, params=None, lr=None, optimizer=None):
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n = 0.0, 0.0, 0
        for X, y in train_iter:
            X, y = X.to(device), y.to(device)
            y_hat = net(X)
            l = loss(y_hat, y)
            # 梯度清零
            if optimizer is None:
                for param in params:
                    if param.grad is not None:
                        param.grad.data.zero_()
            else:
                optimizer.zero_grad()

            l.backward()
            if optimizer is None:
                sgd(params, lr, batch_size)
            else:
                optimizer.step()
            l = l.item()
            train_l_sum += l
            train_acc_sum += (y_hat.argmax(dim=1) == y).sum().item()
            n += y.shape[0]
        test_acc = evaluate_accuracy(test_iter, net)
        print('epoch %d, loss %f, train acc %f, test acc %f'
              % (epoch + 1, train_l_sum / n, train_acc_sum / n, test_acc))

# 训练模型
train_net(net, train_iter, test_iter, loss, num_epochs, batch_size, None, lr, optimizer)

二.VGG

1.模型设计

2.pytorch实现VGG模型

import torch
import numpy as np
import torchvision 
import torchvision.transforms as transforms
from torch import nn

# 超参初始化
lr = 0.05
num_epochs = 10
batch_size = 128

#检查GPU是否可用
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

# 获取MNIST数据集
def load_data_fashion_mnist(batch_size, resize=None):
    # 定义转换，将图像转换为Tensor，并且归一化
    transform = transforms.ToTensor()
    if resize:
        transform = transforms.Compose([
            transforms.Resize(resize),
            transforms.ToTensor()
        ])
    # 加载数据集
    mnist_train = torchvision.datasets.FashionMNIST(root='./database', train=True, download=True, transform=transform)
    mnist_test = torchvision.datasets.FashionMNIST(root='./database', train=False, download=True, transform=transform)

    # 加载数据，windows下只支持num_workers=0，其他可以为4
    train_iter = torch.utils.data.DataLoader(mnist_train, batch_size=batch_size, shuffle=True, num_workers=0)
    test_iter = torch.utils.data.DataLoader(mnist_test, batch_size=batch_size, shuffle=False, num_workers=0)
    return train_iter, test_iter

train_iter, test_iter = load_data_fashion_mnist(batch_size, 224)

# 定义模型
def vgg_block(num_convs, in_channels, out_channels):
    block = []
    for i in range(num_convs):
        block.append(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1))
        block.append(nn.ReLU())
        in_channels = out_channels
    block.append(nn.MaxPool2d(kernel_size=2, stride=2))
    return nn.Sequential(*block)

def vgg(conv_arch):
    conv_block = []
    in_channels = 1
    for (num_convs, out_channels) in conv_arch:
        conv_block.append(vgg_block(num_convs, in_channels, out_channels))
        in_channels = out_channels
    return nn.Sequential(*conv_block,
                        nn.Flatten(),
                        nn.Linear(out_channels * 7 * 7, 4096),
                        nn.ReLU(),
                        nn.Dropout(0.5),
                        nn.Linear(4096, 4096),
                        nn.ReLU(),
                        nn.Dropout(0.5),
                        nn.Linear(4096, 10)
                    )

conv_arch = ((1, 64), (1, 128), (2, 256), (2, 512), (2, 512)) #第一个元素是num_convs,第二个是out_channels
# 降低通道数，方便训练
small_conv_arch = [(pair[0], pair[1]//4) for pair in conv_arch] 
net = vgg(small_conv_arch)
net.to(device)

# 损失函数
loss = nn.CrossEntropyLoss()

# 优化器
optimizer = torch.optim.Adam(net.parameters(), lr=lr)

# 计算准确率
def evaluate_accuracy(data_iter, net):
    acc_sum, n = 0.0, 0
    with torch.no_grad(): # 临时关闭梯度计算
        for X, y in data_iter:
            X, y = X.to(device), y.to(device)
            net.eval() # 评估模式, 关闭dropout
            acc_sum += (net(X).argmax(dim=1) == y).float().sum().item()
            net.train() # 训练模式
            n += y.shape[0]
    return acc_sum / n

# 定义训练模型函数
def train_net(net, train_iter, test_iter, loss, num_epochs, batch_size, params=None, lr=None, optimizer=None):
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n = 0.0, 0.0, 0
        for X, y in train_iter:
            X, y = X.to(device), y.to(device)
            y_hat = net(X)
            l = loss(y_hat, y)
            # 梯度清零
            if optimizer is None:
                for param in params:
                    if param.grad is not None:
                        param.grad.data.zero_()
            else:
                optimizer.zero_grad()

            l.backward()
            if optimizer is None:
                sgd(params, lr, batch_size)
            else:
                optimizer.step()
            l = l.item()
            train_l_sum += l
            train_acc_sum += (y_hat.argmax(dim=1) == y).sum().item()
            n += y.shape[0]
        test_acc = evaluate_accuracy(test_iter, net)
        print('epoch %d, loss %f, train acc %f, test acc %f'
              % (epoch + 1, train_l_sum / n, train_acc_sum / n, test_acc))

# 训练模型
train_net(net, train_iter, test_iter, loss, num_epochs, batch_size, None, lr, optimizer)

三.NiN

1.模型设计

NiN块使用一个普通卷积层和两个1*1卷积层，1*1卷积层卷积层就相当于一个将通道作为特征的全连接层，对每个像素做非线性变换，增强了非线性表达

2.pytorch实现NiN模型

import torch
import numpy as np
import torchvision 
import torchvision.transforms as transforms
from torch import nn

# 超参初始化
lr = 0.1
num_epochs = 10
batch_size = 128

#检查GPU是否可用
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

# 获取MNIST数据集
def load_data_fashion_mnist(batch_size, resize=None):
    # 定义转换，将图像转换为Tensor，并且归一化
    transform = transforms.ToTensor()
    if resize:
        transform = transforms.Compose([
            transforms.Resize(resize),
            transforms.ToTensor()
        ])
    # 加载数据集
    mnist_train = torchvision.datasets.FashionMNIST(root='./database', train=True, download=True, transform=transform)
    mnist_test = torchvision.datasets.FashionMNIST(root='./database', train=False, download=True, transform=transform)

    # 加载数据，windows下只支持num_workers=0，其他可以为4
    train_iter = torch.utils.data.DataLoader(mnist_train, batch_size=batch_size, shuffle=True, num_workers=0)
    test_iter = torch.utils.data.DataLoader(mnist_test, batch_size=batch_size, shuffle=False, num_workers=0)
    return train_iter, test_iter

train_iter, test_iter = load_data_fashion_mnist(batch_size, 224)

# 定义模型
def nin_block(in_channels, out_channels, kernel_size, stride, padding):
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding),
        nn.ReLU(),
        nn.Conv2d(out_channels, out_channels, kernel_size=1),
        nn.ReLU(),
        nn.Conv2d(out_channels, out_channels, kernel_size=1),
        nn.ReLU()
    )

net = nn.Sequential(
    nin_block(1, 96, kernel_size=11, stride=4, padding=0),
    nn.MaxPool2d(kernel_size=3, stride=2),
    nin_block(96, 256, kernel_size=5, stride=1, padding=2),
    nn.MaxPool2d(kernel_size=3, stride=2),
    nin_block(256, 384, kernel_size=3, stride=1, padding=1),
    nn.MaxPool2d(kernel_size=3, stride=2),
    nn.Dropout(0.5),
    nin_block(384, 10, kernel_size=3, stride=1, padding=1),
    nn.AdaptiveAvgPool2d((1, 1)),
    nn.Flatten()
)

net.to(device)

# 损失函数
loss = nn.CrossEntropyLoss()

# 优化器
optimizer = torch.optim.Adam(net.parameters(), lr=lr)

# 计算准确率
def evaluate_accuracy(data_iter, net):
    acc_sum, n = 0.0, 0
    with torch.no_grad(): # 临时关闭梯度计算
        for X, y in data_iter:
            X, y = X.to(device), y.to(device)
            net.eval() # 评估模式, 关闭dropout
            acc_sum += (net(X).argmax(dim=1) == y).float().sum().item()
            net.train() # 训练模式
            n += y.shape[0]
    return acc_sum / n

# 定义训练模型函数
def train_net(net, train_iter, test_iter, loss, num_epochs, batch_size, params=None, lr=None, optimizer=None):
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n = 0.0, 0.0, 0
        for X, y in train_iter:
            X, y = X.to(device), y.to(device)
            y_hat = net(X)
            l = loss(y_hat, y)
            # 梯度清零
            if optimizer is None:
                for param in params:
                    if param.grad is not None:
                        param.grad.data.zero_()
            else:
                optimizer.zero_grad()

            l.backward()
            if optimizer is None:
                sgd(params, lr, batch_size)
            else:
                optimizer.step()
            l = l.item()
            train_l_sum += l
            train_acc_sum += (y_hat.argmax(dim=1) == y).sum().item()
            n += y.shape[0]
        test_acc = evaluate_accuracy(test_iter, net)
        print('epoch %d, loss %f, train acc %f, test acc %f'
              % (epoch + 1, train_l_sum / n, train_acc_sum / n, test_acc))

# 训练模型
train_net(net, train_iter, test_iter, loss, num_epochs, batch_size, None, lr, optimizer)

四.GoogLeNet

1.模型设计

Inception块，有四个通道，对应不同的卷积核，能提升对不同尺寸目标的判别力，白色的1*1卷积层用来降低通道数，减小模型复杂度

2.pytorch实现GoogLeNet模型

import torch
import numpy as np
import torchvision 
import torchvision.transforms as transforms
from torch import nn
from torch.nn import functional as F

# 超参初始化
lr = 0.1
num_epochs = 10
batch_size = 128

#检查GPU是否可用
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

# 获取MNIST数据集
def load_data_fashion_mnist(batch_size, resize=None):
    # 定义转换，将图像转换为Tensor，并且归一化
    transform = transforms.ToTensor()
    if resize:
        transform = transforms.Compose([
            transforms.Resize(resize),
            transforms.ToTensor()
        ])
    # 加载数据集
    mnist_train = torchvision.datasets.FashionMNIST(root='./database', train=True, download=True, transform=transform)
    mnist_test = torchvision.datasets.FashionMNIST(root='./database', train=False, download=True, transform=transform)

    # 加载数据，windows下只支持num_workers=0，其他可以为4
    train_iter = torch.utils.data.DataLoader(mnist_train, batch_size=batch_size, shuffle=True, num_workers=0)
    test_iter = torch.utils.data.DataLoader(mnist_test, batch_size=batch_size, shuffle=False, num_workers=0)
    return train_iter, test_iter

train_iter, test_iter = load_data_fashion_mnist(batch_size, 224)

# 定义模型
class Inception(nn.Module):
    # c1--c4为每条线路里的层的输出通道数
    def __init__(self, in_channels, c1, c2, c3, c4):
        super().__init__()
        # 线路1, 单1x1卷积层
        self.p1_1 = nn.Conv2d(in_channels, c1, kernel_size=1)
        # 线路2, 1x1卷积层后接3x3卷积层
        self.p2_1 = nn.Conv2d(in_channels, c2[0], kernel_size=1)
        self.p2_2 = nn.Conv2d(c2[0], c2[1], kernel_size=3, padding=1)
        # 线路3, 1x1卷积层后接5x5卷积层
        self.p3_1 = nn.Conv2d(in_channels, c3[0], kernel_size=1)
        self.p3_2 = nn.Conv2d(c3[0], c3[1], kernel_size=5, padding=2)
        # 线路4, 3x3最大池化层后接1x1卷积层
        self.p4_1 = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
        self.p4_2 = nn.Conv2d(in_channels, c4, kernel_size=1)

    def forward(self, x):
        p1 = F.relu(self.p1_1(x))
        p2 = F.relu(self.p2_2(F.relu(self.p2_1(x))))
        p3 = F.relu(self.p3_2(F.relu(self.p3_1(x))))
        p4 = F.relu(self.p4_2(self.p4_1(x)))
        # 拼接输出
        return torch.cat((p1, p2, p3, p4), dim=1)

b1 = nn.Sequential(
    nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
)
b2 = nn.Sequential(
    nn.Conv2d(64, 64, kernel_size=1),
    nn.ReLU(),
    nn.Conv2d(64, 192, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
)
b3 = nn.Sequential(
    Inception(192, 64, (96, 128), (16, 32), 32),
    Inception(256, 128, (128, 192), (32, 96), 64),
    nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
)
b4 = nn.Sequential(
    Inception(480, 192, (96, 208), (16, 48), 64),
    Inception(512, 160, (112, 224), (24, 64), 64),
    Inception(512, 128, (128, 256), (24, 64), 64),
    Inception(512, 112, (144, 288), (32, 64), 64),
    Inception(528, 256, (160, 320), (32, 128), 128),
    nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
)
b5 = nn.Sequential(
    Inception(832, 256, (160, 320), (32, 128), 128),
    Inception(832, 384, (192, 384), (48, 128), 128),
    nn.AdaptiveAvgPool2d((1,1)),
    nn.Flatten()
)

net = nn.Sequential(
    b1, b2, b3, b4, b5,
    nn.Linear(1024, 10)
)
net.to(device)

# 损失函数
loss = nn.CrossEntropyLoss()

# 优化器
optimizer = torch.optim.Adam(net.parameters(), lr=lr)

# 计算准确率
def evaluate_accuracy(data_iter, net):
    acc_sum, n = 0.0, 0
    with torch.no_grad(): # 临时关闭梯度计算
        for X, y in data_iter:
            X, y = X.to(device), y.to(device)
            net.eval() # 评估模式, 关闭dropout
            acc_sum += (net(X).argmax(dim=1) == y).float().sum().item()
            net.train() # 训练模式
            n += y.shape[0]
    return acc_sum / n

# 定义训练模型函数
def train_net(net, train_iter, test_iter, loss, num_epochs, batch_size, params=None, lr=None, optimizer=None):
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n = 0.0, 0.0, 0
        for X, y in train_iter:
            X, y = X.to(device), y.to(device)
            y_hat = net(X)
            l = loss(y_hat, y)
            # 梯度清零
            if optimizer is None:
                for param in params:
                    if param.grad is not None:
                        param.grad.data.zero_()
            else:
                optimizer.zero_grad()

            l.backward()
            if optimizer is None:
                sgd(params, lr, batch_size)
            else:
                optimizer.step()
            l = l.item()
            train_l_sum += l
            train_acc_sum += (y_hat.argmax(dim=1) == y).sum().item()
            n += y.shape[0]
        test_acc = evaluate_accuracy(test_iter, net)
        print('epoch %d, loss %f, train acc %f, test acc %f'
              % (epoch + 1, train_l_sum / n, train_acc_sum / n, test_acc))

# 训练模型
train_net(net, train_iter, test_iter, loss, num_epochs, batch_size, None, lr, optimizer)

五.ResNet

1.模型设计

残差块将目标映射从f(x)转换到f(x)-x，降低了优化难度，缓解了梯度消失的问题，并使损失曲面更加平滑

包含1*1卷积层的残差块，主要是用来调整通道和分辨率

2.pytorch实现ResNet模型

import torch
import numpy as np
import torchvision 
import torchvision.transforms as transforms
from torch import nn
from torch.nn import functional as F

# 超参初始化
lr = 0.05
num_epochs = 10
batch_size = 256

#检查GPU是否可用
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

# 获取MNIST数据集
def load_data_fashion_mnist(batch_size, resize=None):
    # 定义转换，将图像转换为Tensor，并且归一化
    transform = transforms.ToTensor()
    if resize:
        transform = transforms.Compose([
            transforms.Resize(resize),
            transforms.ToTensor()
        ])
    # 加载数据集
    mnist_train = torchvision.datasets.FashionMNIST(root='./database', train=True, download=True, transform=transform)
    mnist_test = torchvision.datasets.FashionMNIST(root='./database', train=False, download=True, transform=transform)

    # 加载数据，windows下只支持num_workers=0，其他可以为4
    train_iter = torch.utils.data.DataLoader(mnist_train, batch_size=batch_size, shuffle=True, num_workers=0)
    test_iter = torch.utils.data.DataLoader(mnist_test, batch_size=batch_size, shuffle=False, num_workers=0)
    return train_iter, test_iter

train_iter, test_iter = load_data_fashion_mnist(batch_size, 224)

# 定义模型
class Residual(nn.Module):
    """
    残差块
    """
    def __init__(self, in_channels, out_channels, use_1x1conv=False, strides=1):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, stride=strides)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        if use_1x1conv:
            self.conv3 = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=strides)
        else:
            self.conv3 = None
        # 批量归一化
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, X):
        Y = self.relu(self.bn1(self.conv1(X)))
        Y = self.bn2(self.conv2(Y))
        if self.conv3:
            X = self.conv3(X)
        return self.relu(Y + X)

# 多个残差块组合
def resnet_block(in_channels, out_channels, num_residuals, first_block=False):
    blk = []
    for i in range(num_residuals):
        if i == 0 and not first_block:
            blk.append(Residual(in_channels, out_channels, use_1x1conv=True, strides=2))
        else:
            blk.append(Residual(out_channels, out_channels))
    return nn.Sequential(*blk)

b1 = nn.Sequential(
    nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3),
    nn.BatchNorm2d(64),
    nn.ReLU(inplace=True),
    nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
)


net = nn.Sequential(
    b1,
    resnet_block(64, 64, 2, first_block=True),
    resnet_block(64, 64, 2),
    resnet_block(64, 128, 2),
    resnet_block(128, 256, 2),
    resnet_block(256, 512, 2),
    nn.AdaptiveAvgPool2d((1, 1)),
    nn.Flatten(),
    nn.Linear(512, 10)
)

net.to(device)

# 损失函数
loss = nn.CrossEntropyLoss()

# 优化器
optimizer = torch.optim.Adam(net.parameters(), lr=lr)

# 计算准确率
def evaluate_accuracy(data_iter, net):
    acc_sum, n = 0.0, 0
    with torch.no_grad(): # 临时关闭梯度计算
        for X, y in data_iter:
            X, y = X.to(device), y.to(device)
            net.eval() # 评估模式, 关闭dropout
            acc_sum += (net(X).argmax(dim=1) == y).float().sum().item()
            net.train() # 训练模式
            n += y.shape[0]
    return acc_sum / n

# 定义训练模型函数
def train_net(net, train_iter, test_iter, loss, num_epochs, batch_size, params=None, lr=None, optimizer=None):
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n = 0.0, 0.0, 0
        for X, y in train_iter:
            X, y = X.to(device), y.to(device)
            y_hat = net(X)
            l = loss(y_hat, y)
            # 梯度清零
            if optimizer is None:
                for param in params:
                    if param.grad is not None:
                        param.grad.data.zero_()
            else:
                optimizer.zero_grad()

            l.backward()
            if optimizer is None:
                sgd(params, lr, batch_size)
            else:
                optimizer.step()
            l = l.item()
            train_l_sum += l
            train_acc_sum += (y_hat.argmax(dim=1) == y).sum().item()
            n += y.shape[0]
        test_acc = evaluate_accuracy(test_iter, net)
        print('epoch %d, loss %f, train acc %f, test acc %f'
              % (epoch + 1, train_l_sum / n, train_acc_sum / n, test_acc))

# 训练模型
train_net(net, train_iter, test_iter, loss, num_epochs, batch_size, None, lr, optimizer)

六.DenseNet

1.模型设计

DenseNet是ResNet的改良，将ResNet的相加改为连结，即保持空间尺度不变的情况下，增加通道数，使每一层都能直接利用前面所有层的特征，使梯度更容易回传，并减小了总参数

2.pytorch实现DenseNet模型

import torch
import numpy as np
import torchvision 
import torchvision.transforms as transforms
from torch import nn
from torch.nn import functional as F

# 超参初始化
lr = 0.05
num_epochs = 10
batch_size = 256

#检查GPU是否可用
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

# 获取MNIST数据集
def load_data_fashion_mnist(batch_size, resize=None):
    # 定义转换，将图像转换为Tensor，并且归一化
    transform = transforms.ToTensor()
    if resize:
        transform = transforms.Compose([
            transforms.Resize(resize),
            transforms.ToTensor()
        ])
    # 加载数据集
    mnist_train = torchvision.datasets.FashionMNIST(root='./database', train=True, download=True, transform=transform)
    mnist_test = torchvision.datasets.FashionMNIST(root='./database', train=False, download=True, transform=transform)

    # 加载数据，windows下只支持num_workers=0，其他可以为4
    train_iter = torch.utils.data.DataLoader(mnist_train, batch_size=batch_size, shuffle=True, num_workers=0)
    test_iter = torch.utils.data.DataLoader(mnist_test, batch_size=batch_size, shuffle=False, num_workers=0)
    return train_iter, test_iter

train_iter, test_iter = load_data_fashion_mnist(batch_size, 224)

# 定义模型
def conv_block(in_channels, out_channels):
    return nn.Sequential(
        nn.BatchNorm2d(in_channels),
        nn.ReLU(),
        nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
    )

class DenseBlock(nn.Module):
    """
    稠密块
    """
    def __init__(self, num_convs, in_channels, num_channels):
        super().__init__()
        block = []
        for i in range(num_convs):
            block.append(conv_block(num_channels * i + in_channels, num_channels))
        self.net = nn.Sequential(*block)

    def forward(self, X):
        for block in self.net:
            Y = block(X)
            X = torch.cat((X, Y), dim=1)
        return X

# 过渡层，用1*1卷积层减小通道数
def transition_block(in_channels, out_channels):
    return nn.Sequential(
        nn.BatchNorm2d(in_channels),
        nn.ReLU(),
        nn.Conv2d(in_channels, out_channels, kernel_size=1),
        nn.AvgPool2d(kernel_size=2, stride=2)
    )

b1 = nn.Sequential(
    nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3),
    nn.BatchNorm2d(64),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
)

num_channels, growth_rate = 64, 32
num_convs_in_dense_blocks = [4, 4, 4, 4] # 每个稠密块的卷积层数量
b2 = []

for i, num_convs in enumerate(num_convs_in_dense_blocks):
    b2.append(DenseBlock(num_convs, num_channels, growth_rate))
    num_channels += num_convs * growth_rate
    # 在稠密块后加使通道减半的过渡层
    if i != len(num_convs_in_dense_blocks) - 1:
        b2.append(transition_block(num_channels, num_channels // 2))
        num_channels = num_channels // 2

net = nn.Sequential(
    b1,
    *b2,
    nn.BatchNorm2d(num_channels),
    nn.ReLU(),
    nn.AdaptiveAvgPool2d((1, 1)),
    nn.Flatten(),
    nn.Linear(num_channels, 10)
)

net.to(device)

# 损失函数
loss = nn.CrossEntropyLoss()

# 优化器
optimizer = torch.optim.Adam(net.parameters(), lr=lr)

# 计算准确率
def evaluate_accuracy(data_iter, net):
    acc_sum, n = 0.0, 0
    with torch.no_grad(): # 临时关闭梯度计算
        for X, y in data_iter:
            X, y = X.to(device), y.to(device)
            net.eval() # 评估模式, 关闭dropout
            acc_sum += (net(X).argmax(dim=1) == y).float().sum().item()
            net.train() # 训练模式
            n += y.shape[0]
    return acc_sum / n

# 定义训练模型函数
def train_net(net, train_iter, test_iter, loss, num_epochs, batch_size, params=None, lr=None, optimizer=None):
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n = 0.0, 0.0, 0
        for X, y in train_iter:
            X, y = X.to(device), y.to(device)
            y_hat = net(X)
            l = loss(y_hat, y)
            # 梯度清零
            if optimizer is None:
                for param in params:
                    if param.grad is not None:
                        param.grad.data.zero_()
            else:
                optimizer.zero_grad()

            l.backward()
            if optimizer is None:
                sgd(params, lr, batch_size)
            else:
                optimizer.step()
            l = l.item()
            train_l_sum += l
            train_acc_sum += (y_hat.argmax(dim=1) == y).sum().item()
            n += y.shape[0]
        test_acc = evaluate_accuracy(test_iter, net)
        print('epoch %d, loss %f, train acc %f, test acc %f'
              % (epoch + 1, train_l_sum / n, train_acc_sum / n, test_acc))

# 训练模型
train_net(net, train_iter, test_iter, loss, num_epochs, batch_size, None, lr, optimizer)