torch，MNIST实践

import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np

# 定义神经网络
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 卷积层Conv2d输入通道数1,输出通道数5,卷积核大小6x6。输入图像大小为(H,W),
				 # 经过conv1输出特征图大小为(H-5,W-5)。
        self.conv1 = torch.nn.Conv2d(1, 5, 6)
        # 第二卷积层Conv2d的参数个第一个参数是输入的通道数量是上一个输出数量5，第二个是输出的通道数量8，第三个是神经核数量5*5
        self.conv2 = torch.nn.Conv2d(5, 8, 5)
        # 池化层 池化核大小为2x2,步长为2。特征图大小变为(H-5)/2,(W-5)/2。每一次卷积层都要处理一次
        self.pool = torch.nn.MaxPool2d(2, 2)
        # 由于上一层有16个channel输出，每个feature map大小为3*3，所以全连接层的输入是8*3*3，
				 # （(原始数据 H-第一次卷积核数量)/2-第二次卷积核数量）/2 在四舍五入 * （(原始数据 W-第一次卷积核数量)/2-第二次卷积核数量）/2 在四舍五入
        self.fc1 = torch.nn.Linear(8*3*3, 120)
        self.fc2 = torch.nn.Linear(120, 84)
        self.fc3 = torch.nn.Linear(84, 10)
        # 最终有10类，所以最后一个全连接层输出数量是10

    def forward(self, x):
        x = self.pool(torch.nn.functional.relu(self.conv1(x)))
        x = self.pool(torch.nn.functional.relu(self.conv2(x)))
        x = x.view(-1, 8*3*3) #第二个参数数据要和上面的fc1 的第一个参数一致
        x = torch.nn.functional.relu(self.fc1(x))
        x = torch.nn.functional.relu(self.fc2(x))
        x = self.fc3(x)
        return x


def Training():
    net = Net()
    # 定义数据预处理
    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5,), (0.5,))])

    # 加载数据集
    trainset = torchvision.datasets.MNIST(root='./data', train=True,
                                            download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=1,
                                              shuffle=True,)

    # 定义损失函数和优化器
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

    # 训练网络
    for epoch in range(1):  # 多次循环遍历数据集
        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data
            # 首先要把梯度清零，不然PyTorch每次计算梯度会累加，不清零的话第二次算的梯度等于第一次加第二次的
            optimizer.zero_grad()
            # 计算前向传播的输出
            outputs = net(inputs)
            # 根据输出计算loss
            loss = criterion(outputs, labels)
            # print(loss,labels)
            # 算完loss之后进行反向梯度传播，这个过程之后梯度会记录在变量中
            loss.backward()
            # 用计算的梯度去做优化
            optimizer.step()
            running_loss += loss.item()
            if i % 2000 == 1999:    # 每 2000 批次打印一次平均损失值
                print('[%d, %5d] loss: %.3f' %
                      (epoch + 1, i + 1, running_loss / 2000))
                running_loss = 0.0
                # show_images(inputs,labels)
            # print(outputs)
            # break
    # 保存模型
    torch.save(net.state_dict(), 'mnist_cnn.pth')
    print('Finished Training')

def show_images(images,labels):
    img = torchvision.utils.make_grid(images)
    img = img.numpy().transpose(1, 2, 0)
    std = [0.5, 0.5, 0.5]
    mean = [0.5, 0.5, 0.5]
    img = img * std + mean
    print(labels)
    plt.imshow(img)
    plt.show()

def test_load(path):
    img = Image.open(path)
    img = img.convert('L')  # 转为灰度图
    img = img.resize((28, 28))  # 改为28x28
    img = np.array(img)
    img = (img / 255 - 0.5) / 0.5  # 标准化数据值在-1到1
    img = img[np.newaxis, np.newaxis, :, :]  # 增加batch和channel维度
    img = torch.Tensor(img)
    return img

#训练和测试一体
def custom_training(img_path,result=None):
    model = Net()
    model.load_state_dict(torch.load('mnist_cnn.pth'))
    # 切换到预测模式
    model.eval()
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
    img = test_load(img_path)
    outputs = model(img)
    _, predicted = torch.max(outputs.data, 1)
    if result != None:
        labels = torch.tensor([result],dtype=torch.int64)
        # 根据输出计算loss
        l = 1
        while l !=0:
            optimizer.zero_grad()
            # 计算前向传播的输出
            outputs = model(img)
            flattened_tensor = outputs.view(-1)
            # 找到最大值的索引
            max_value_index = torch.argmax(flattened_tensor)
            if max_value_index == result:
                break
            loss = criterion(outputs, labels)
            # 算完loss之后进行反向梯度传播，这个过程之后梯度会记录在变量中
            loss.backward()
            # 用计算的梯度去做优化
            optimizer.step()
            l=loss
					#保存模型
        torch.save(model.state_dict(), 'mnist_cnn.pth')
        print('Finished Training',optimizer)
    else:
        print(outputs)
        print(predicted)


Training()
# test_model()
# test_custom_model('test_3.png')
# custom_training('test9.png'