pytorch 基本CNN模型使用

import torch
from torchvision import datasets
from torchvision import transforms
from torch.utils.data import DataLoader
import torch.nn.functional as F
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
 
# 数据正态化
batch_size = 64
loss_list = []
accuracy_list = []
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307, ), (0.3081, ))
])
# 加载测试与训练数据集
train_dataset = datasets.MNIST(root='dataset/mnist/',
               train=True,
               download=True,
               transform = transform)
 
train_load = DataLoader(train_dataset,
                        batch_size=64,
                        shuffle=True,
                        num_workers=6)
 
test_dataset = datasets.MNIST(root='dataset/mnist/',
                              train=False,
                              download=True,
                              transform= transform)
 
test_load = DataLoader(test_dataset,
                       batch_size=batch_size,
                       shuffle=True,
                       num_workers=6)
 
# 网络模型
class Net(torch.nn.Module):
    """
    1*10*5*5 的卷积层 + relu激活函数 + 最大池化层2*2
    10*20*5*5 的卷积层 + relu激活函数 + 最大池化层2*2 
    将20*4*4 转换为一维张量320*1
    然后全连接层输出10*1
    """
    def __init__(self):
        super(Net, self).__init__()
        self.conv_layer1 = torch.nn.Conv2d(in_channels=1, out_channels=10, kernel_size=5)
        self.conv_layer2 = torch.nn.Conv2d(in_channels=10, out_channels=20, kernel_size=5)
        self.pool_layer = torch.nn.MaxPool2d(kernel_size=2)
        self.linear_layer = torch.nn.Linear(320, 10)
 
    def forward(self, x):
        batch_size = x.size()[0]
        x = F.relu(self.pool_layer(self.conv_layer1(x)))
        x = F.relu(self.pool_layer(self.conv_layer2(x)))
        x = x.view(batch_size, -1)
        x = self.linear_layer(x)
        return x
 
model = Net()
# GPU 加速
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
criterion = nn.CrossEntropyLoss(reduction='mean')
optimizer = optim.SGD(model.parameters(), lr=0.009, momentum=0.5)
 
 
 
def train(epoch):
    run_loss = 0.0
    for num, data in enumerate(train_load, 0):
        inputs, labels = data
        # print(inputs.shape)
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        run_loss += loss.item()
 
        if(num % 300==299):
            print("epoch=%d,num=%d, loss=%.3f" %(epoch+1, num+1, run_loss/300))
            loss_list.append(run_loss)
            run_loss = 0
 
def run_data():
    total = 0
    correct = 0
    with torch.no_grad():
        for data in test_load:
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, dim=1)
            correct += (predicted == labels).sum().item()
            total += labels.size()[0]
    accuracy = 100 * correct / total
    accuracy_list.append(accuracy)
    print("Accuracy on test set :%.2f %% [%d, %d]" %(accuracy, correct, total))
 
if __name__ == '__main__':
    # run_data()
    for epoch in range(30):
        train(epoch)
        run_data()
    plt.plot(loss_list)
    plt.show()
    plt.plot(accuracy_list)
    plt.show()