多层全连接网络实现MNIST手写数字分类

# 多层全连接网络实现MNIST手写数字分类
import numpy as np
import torch
import torchvision
from torch import nn, optim
from torch.utils.data import DataLoader
from torch.optim import SGD
from torchvision.transforms import transforms
from torch.autograd.variable import Variable


# 简单的三层全连接神经网络
class Simple_Net(nn.Module):
    def __init__(self, in_dim, n_hidden_1, n_hidden_2, out_dim):
        super(Simple_Net, self).__init__()
        self.layer1 = nn.Linear(in_dim, n_hidden_1)  # 第一个全连接层，将输入纬度(in_dim)转换为第一隐藏层纬度(n_hidden_1)
        self.layer2 = nn.Linear(n_hidden_1, n_hidden_2)
        self.layer3 = nn.Linear(n_hidden_2, out_dim)

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        return x


# 添加激活函数，增加网络的非线性
class Activation_Net(nn.Module):
    def __init__(self, in_dim, n_hidden_1, n_hidden_2, out_dim):
        super(Activation_Net, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Linear(in_dim, n_hidden_1),
            nn.ReLU(True)
        )
        self.layer2 = nn.Sequential(
            nn.Linear(n_hidden_1, n_hidden_2),
            nn.ReLU(True)
        )
        self.layer3 = nn.Linear(n_hidden_2, out_dim)

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        return x


# 添加批标准化
class Batch_Net(nn.Module):
    def __init__(self, in_dim, n_hidden_1, n_hidden_2, out_dim):
        super(Batch_Net, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Linear(in_dim, n_hidden_1),
            nn.BatchNorm1d(n_hidden_1),
            nn.ReLU(True)
        )
        self.layer2 = nn.Sequential(
            nn.Linear(n_hidden_1, n_hidden_2),
            nn.BatchNorm1d(n_hidden_2),  # 批归一化层
            nn.ReLU(True)
        )
        self.layer3 = nn.Linear(n_hidden_2, out_dim)

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        return x


# 超参数
batch_size = 64
learning_rate = 1e-2
num_epochs = 20

# 数据预处理
data_tf = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize([0.5], [0.5])]
)

# 获取数据集
train_dataset = torchvision.datasets.MNIST("data", train=True, transform=data_tf, download=True)
test_dataset = torchvision.datasets.MNIST("data", train=False, transform=data_tf, download=True)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 初始化参数
D_in, H1, H2, D_out = 28, 300, 100, 10

model = Simple_Net(D_in * D_in, H1, H2, D_out)

if torch.cuda.is_available():  # 模型实例化并移到GPU设备
    model = model.cuda()

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)


# 开始训练网络
def train_model(model, criterion, optimizer, num_epochs):
    for epoch in range(num_epochs):
        print('epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        train_loss = 0.0
        train_acc = 0.0
        # 获取数据输入和标签，封装成变量
        for data in train_loader:  # 获得一个batch样本
            img, label = data  # 获得图片和标签
            img = img.view(img.size(0), -1)  # 将图片从(64（batch_size）,1,28,28)转换为（64，28*28）
            if torch.cuda.is_available():
                img = Variable(img).cuda()
                label = Variable(label).cuda()
            else:
                img, label = Variable(img), Variable(label)
            # 梯度参数清零
            optimizer.zero_grad()
            # 前向
            out = model(img)  # 等价于 out = model.forward(img)
            loss = criterion(out, label)
            _, preds = torch.max(out.data, 1)
            # 反向传播
            loss.backward()
            optimizer.step()
            # 统计
            train_loss += loss.item()
            train_correct = torch.sum(preds == label.data)
            train_acc += train_correct
        print(
            'Train Loss: {:.6f}, Acc: {:.6f}'.format(train_loss / (len(train_dataset)), train_acc / (len(train_dataset))))

        model.eval()  # 设置模型为评估模式
        eval_loss = 0.0
        eval_acc = 0.0
        for data in test_loader:
            img, label = data
            img = img.view(img.size(0), -1)
            if torch.cuda.is_available():
                with torch.no_grad():
                    img = Variable(img).cuda()
                    label = Variable(label).cuda()
            else:
                img = Variable(img, volatile=True)
                label = Variable(label, volatile=True)
            out = model(img)
            loss = criterion(out, label)
            eval_loss += loss.item()
            _, preds = torch.max(out.data, 1)
            num_correct = torch.sum(preds == label.data)
            eval_acc += num_correct
        print('Test Loss:{:.6f}, Acc: {:.6f}'.format(eval_loss / (len(test_dataset)), eval_acc / (len(test_dataset))))


print(train_model(model, criterion, optimizer, num_epochs))