数字识别神经网络搭建

网络模型搭建

import numpy
	import torch
	# 数字识别模型建立
	class Net(torch.nn.Module):
	    def __init__(self):
	        super(Net, self).__init__()

        self.fc = torch.nn.Sequential(
            # 拉平，将 (28, 28) => (784) 1维数组
            torch.nn.Flatten(),
            torch.nn.Linear(28*28, 24*24),
            torch.nn.ReLU(),
            torch.nn.Linear(24*24, 12*12),
            torch.nn.ReLU(),
            torch.nn.Linear(12*12, 10),
            torch.nn.ReLU(),
        )

        # 优化器定义
        self.optim = torch.optim.SGD(self.parameters(), lr=0.01)
        # 损失函数定义
        self.loss = torch.nn.CrossEntropyLoss()

    def forward(self, x):
        # 模型执行过程
        result = self.fc(x)
        return result

    def back(self, x):
        # 反向传播
        y = numpy.NaN
        # 梯度操作
        loss = self.loss(x, y)
        # 梯度更新
        self.optim.step()
        print("loss:", loss.data.numpy())
        pass

模型训练

import time
import torch
import torchvision
import torch.utils.data as Data
import net_moudle


# 网络模型初始化
def _createNetModle():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    net = net_moudle.Net()
    net = net.to(device)
    # 网络结构输出
    print(net)
    return net

# 网络模型加载
def _loadNetModle():
    net = net_moudle.Net()
    net.to()
    net.load_state_dict(torch.load("./data/net_param_20230914134640_0.00017671784735284746.plk"))
    # 网络结构输出
    print(net)
    return net


# 网络模型训练
def _train(net: net_moudle.Net):
    for step, (batch_x, batch_y) in enumerate(train_loader):
        # 前向传播结构获取
        pred_y = net(batch_x)
        # 学习情况
        loss = _back(net, pred_y, batch_y)
        # if step % 50 == 0:
    return loss


# 反向传播
def _back(net: net_moudle.Net, pred_y, batch_y):
    loss = net.loss(pred_y, batch_y)
    # 清空上一层梯度
    net.optim.zero_grad()
    # 反向传播
    loss.backward()
    # 步进
    net.optim.step()
    return loss


# 模型保存
def _save(net: net_moudle.Net, loss):
    root = "./data/{}_{}_{}.plk".format(
        "net",
        time.strftime("%Y%m%d%H%M%S", time.localtime(int(time.time()))),
        loss.data.numpy(),
    )
    # 网络保存
    torch.save(net, root)

    root = "./data/{}_{}_{}.plk".format(
        "net_param",
        time.strftime("%Y%m%d%H%M%S", time.localtime(int(time.time()))),
        loss.data.numpy(),
    )
    # 参数保存
    torch.save(net.state_dict(), root)


# 训练集
train_data = torchvision.datasets.MNIST(
    root='./MNIST/',
    train=True,
    transform=torchvision.transforms.ToTensor(),
    download=True
)

# 测试集
test_data = torchvision.datasets.MNIST(root='./MNIST/', train=False)  # train设置为False表示获取测试集

# 一个批训练 50个样本, 1 channel通道, 图片尺寸 28x28 size:(50, 1, 28, 28)
train_loader = Data.DataLoader(
    dataset=train_data,
    batch_size=50,
    shuffle=True
)

#  测试数据预处理；只测试前2000个
test_x = torch.unsqueeze(test_data.data, dim=1).float()[:2000] / 255.0
# shape from (2000, 28, 28) to (2000, 1, 28, 28)
test_y = test_data.targets[:2000]

# 记录做最好的的模型
list = [0.00017671784735284746]

# 设置随机种子, 用于复现
torch.manual_seed(1)

if __name__ == '__main__':
    # net = _createNetModle()
    net = _loadNetModle()
    for epoch in range(100):
        loss = _train(net)
        if min(list) > loss.data.numpy():
            list.append(loss.data.numpy())
            _save(net, loss)
        print(f'Epoch [{epoch + 1}/{100}] train loss: {loss.data.numpy():.6f}')

    for p in net.parameters():
        print(p.shape, p.numel())
    test_output = net(test_x[:30])
    pred_y = torch.max(test_output, 1)[1].numpy()
    print(pred_y, 'prediction number')
    print(test_y[:30].numpy(), 'real number')