mnist模型训练（附代码）

获取各种模型的github地址1：https://github.com/Att100/CIFAR10_Pytorch.
获取各种模型的github地址2：https://github.com/chenyaofo/pytorch-cifar-models.

训练演示

环境：

python=3.6 ~ 3.8
Package                 Version
----------------------- ---------
torch                   1.10.1
torchvision             0.11.2
tqdm                    4.62.3

文件目录展示：

根目录:

这个data文件是运行训练代码时候会自动创建并下载数据的。
model中的文件是下面AlexNet的py文件

模型代码：AlexNet.py

import torch
import torch.nn as nn
import torch.nn.functional as F


class AlexNet(nn.Module):
    def __init__(self):
        super(AlexNet, self).__init__()
        self.Conv2d_1 = nn.Conv2d(kernel_size=3, in_channels=1, out_channels=96, padding=1)
        self.bn_1 = nn.BatchNorm2d(96)
        self.maxpool_1 = nn.MaxPool2d((3, 3), stride=2, padding=1)

        self.Conv2d_2 = nn.Conv2d(kernel_size=5, in_channels=96, out_channels=256, padding=2)
        self.bn_2 = nn.BatchNorm2d(256)
        self.maxpool_2 = nn.MaxPool2d((3, 3), stride=2, padding=1)

        self.Conv2d_3 = nn.Conv2d(kernel_size=3, in_channels=256, out_channels=384, padding=1)
        self.Conv2d_4 = nn.Conv2d(kernel_size=3, in_channels=384, out_channels=384, padding=1)
        self.Conv2d_5 = nn.Conv2d(kernel_size=3, in_channels=384, out_channels=256, padding=1)
        self.bn_3 = nn.BatchNorm2d(256)
        self.maxpool_3 = nn.MaxPool2d((3, 3), stride=2, padding=1)

        self.fc_1 = nn.Linear(4*4*256, 2048)
        self.dp_1 = nn.Dropout()
        self.fc_2 = nn.Linear(2048, 1024)
        self.dp_2 = nn.Dropout()
        self.fc_3 = nn.Linear(1024, 10)

    def forward(self, x):
        x = self.Conv2d_1(x)
        x = self.bn_1(x)
        x = F.relu(x)
        x = self.maxpool_1(x)

        x = self.Conv2d_2(x)
        x = self.bn_2(x)
        x = F.relu(x)
        x = self.maxpool_2(x)

        x = F.relu(self.Conv2d_3(x))
        x = F.relu(self.Conv2d_4(x))
        x = F.relu(self.Conv2d_5(x))
        x = self.bn_3(x)
        x = F.relu(x)
        x = self.maxpool_3(x)

        x = x.view(-1, 4*4*256)
        x = F.relu(self.fc_1(x))
        x = self.dp_1(x)
        x = F.relu(self.fc_2(x))
        x = self.dp_2(x)
        x = self.fc_3(x)
        return x

训练代码：train_alexnet.py

import torch
import codecs
import itertools
import torch.nn.init as init
import torch.nn as nn
from tqdm import tqdm
from model import AlexNet
import torchvision.utils as utils
import torch.optim as optim
from torchvision import datasets
from torch.autograd import Variable
from torch.utils.data import DataLoader
import torchvision.transforms as transforms


transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])

train_mnist = DataLoader(datasets.MNIST('data', train=True, download=True, transform=transform), batch_size=128, shuffle=True)

alexNet = AlexNet.AlexNet()


learning_rate = 0.001
momentum = 0.9
epoches = 200

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

alexNet = alexNet.to(device)

# optimizer and loss function
optimizer = optim.SGD(alexNet.parameters(), lr=learning_rate, momentum=momentum, weight_decay=1e-5)
loss_func = torch.nn.CrossEntropyLoss()

for epoch in range(epoches):
	train_mnist = tqdm(train_mnist)
	running_loss = 0.0
	for inputs, labels in train_mnist:
		inputs,labels = inputs.to(device),labels.to(device)
		# ============= forward =============
		outputs = alexNet(inputs)
		# ============= backward ============
		optimizer.zero_grad()
		loss = loss_func(outputs, labels)
		loss.backward()
		optimizer.step()
		# ============= log information =====
		running_loss += loss.item()
		description = 'epoch: %d , current_loss: %.4f, running_loss: %.4f' % (epoch, loss.item(), running_loss)
		train_mnist.set_description(description)
		train_mnist.update()

	torch.save(alexNet.state_dict(), './checkpoint/checkpoint_' + str(epoch) + '.pt')

测试模型泛化性能代码: test_model_accuracy.py

import torch
import codecs
import itertools
import torch.nn.init as init
import torch.nn as nn
from tqdm import tqdm
from model import AlexNet
import torchvision.utils as utils
import torch.optim as optim
from torchvision import datasets
from torch.autograd import Variable
from torch.utils.data import DataLoader
import torchvision.transforms as transforms

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])

train_mnist = DataLoader(datasets.MNIST('data', train=False, download=True, transform=transform), batch_size=128, shuffle=True)

alexNet = AlexNet.AlexNet()

alexNet.load_state_dict(torch.load("./checkpoint/checkpoint_10.pt", map_location=device))



alexNet = alexNet.to(device)
alexNet = alexNet.eval()

with torch.no_grad():
    train_mnist = tqdm(train_mnist)
    total = 0
    correct = 0
    for inputs, labels in train_mnist:
            inputs,labels = inputs.to(device),labels.to(device)
            # ============= forward =============
            outputs = alexNet(inputs)
            # ============= precision ===========
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            description = 'correct: %.4f, total: %.4f , accuracy: %.4f' % (correct, total, correct/total)
            train_mnist.set_description(description)
            train_mnist.update()
            
            

😮超级✔海量✔视觉✔模型😉：https://github.com/52CV/CVPR-2021-Papers.