跟着李沐学AI代码复现-03BN-LeNet,04ResNet

最新推荐文章于 2024-08-04 10:55:56 发布

合乎周礼

最新推荐文章于 2024-08-04 10:55:56 发布

阅读量92

点赞数 1

文章标签：人工智能 python 深度学习

本文链接：https://blog.csdn.net/m0_55895748/article/details/140401747

版权

import torch
from torch import nn
from torch.utils import data

import torchvision
from torchvision import transforms

'''
'''

# 超参数定义
batch_size = 256
epoch_num = 10
lr = 0.03
if torch.cuda.device_count() >= 1:
    device = torch.device(f'cuda:0')
else:
    device = torch.device('cpu')


def load():
    # 加载数据集
    # 这里的root指的是本地目录。
    mnist_train = torchvision.datasets.FashionMNIST(
        root="data", train=True, transform=transforms.ToTensor(), download=True)
    mnist_test = torchvision.datasets.FashionMNIST(
        root="data", train=False, transform=transforms.ToTensor(), download=True)
    # 将数据集转为特定格式 
    train_iter = data.DataLoader(mnist_train, batch_size, shuffle=True, num_workers=8)
    test_iter = data.DataLoader(mnist_test, batch_size, shuffle=False, num_workers=8)

    return train_iter, test_iter


def netdefine():
    # 网络定义
    net = nn.Sequential(
        nn.Conv2d(1, 6, kernel_size=5, padding=2), nn.BatchNorm2d(6),
        nn.Sigmoid(), nn.AvgPool2d(kernel_size=2, stride=2),
        nn.Conv2d(6, 16, kernel_size=5), nn.BatchNorm2d(16),
        nn.Sigmoid(),
        nn.AvgPool2d(kernel_size=2, stride=2),
        nn.Flatten(),
        nn.Linear(16 * 5 * 5, 120), nn.BatchNorm1d(120),
        nn.Sigmoid(),
        nn.Linear(120, 84), nn.BatchNorm1d(84),
        nn.Sigmoid(),
        nn.Linear(84, 10))

    # 初始化参数
    def init_weights(m):
        if type(m) == nn.Linear or type(m) == nn.Conv2d:
            nn.init.xavier_uniform_(m.weight)

    net.apply(init_weights)
    # 模型迁移
    net = net.to(device)
    return net


def train(net, train_iter):
    # optimizer会根据梯度对所有参数进行更新
    optimizer = torch.optim.SGD(net.parameters(), lr)
    loss = nn.CrossEntropyLoss()
    for epoch in range(0, epoch_num):
        net.train()
        total = 0
        correct = 0
        for X, y in train_iter:
            # 梯度置零
            X, y = X.to(device), y.to(device)
            y_hat = net(X)
            l = loss(y_hat, y)
            # 梯度反向传播，参数更新
            optimizer.zero_grad()
            l.backward()
            optimizer.step()
            # 计算准确率
            _, pred = torch.max(y_hat, 1)
            total += y.size(0)
            correct += (pred == y).sum().item()
        print(f'epoch = {epoch}, total = {total}, acurrcy = {100 * correct / total}%')


def test(net, test_iter):
    net.eval()
    with torch.no_grad():
        correct = 0
        total = 0
        for X, y in test_iter:
            X = X.to(device)
            y = y.to(device)
            output = net(X)
            _, pred = torch.max(output, 1)
            total += y.size(0)
            correct += (pred == y).sum().item()
        print(f'test,total = {total}, acurrcy = {100 * correct / total}%')


if __name__ == '__main__':
    train_iter, test_iter = load()
    net = netdefine()
    train(net, train_iter)
    test(net, test_iter)

import torch
from torch import nn
from torch.utils import data

import torchvision
from torchvision import transforms

'''
'''

# 超参数定义
batch_size = 256
epoch_num = 10
lr = 0.03
if torch.cuda.device_count() >= 1:
    device = torch.device(f'cuda:0')
else:
    device = torch.device('cpu')


def load():
    # 加载数据集
    # 这里的root指的是本地目录。
    mnist_train = torchvision.datasets.FashionMNIST(
        root="data", train=True, transform=transforms.ToTensor(), download=True)
    mnist_test = torchvision.datasets.FashionMNIST(
        root="data", train=False, transform=transforms.ToTensor(), download=True)
    # 将数据集转为特定格式 
    train_iter = data.DataLoader(mnist_train, batch_size, shuffle=True, num_workers=8)
    test_iter = data.DataLoader(mnist_test, batch_size, shuffle=False, num_workers=8)

    return train_iter, test_iter


def netdefine():
    # 网络定义
    net = nn.Sequential(
        nn.Conv2d(1, 6, kernel_size=5, padding=2), nn.BatchNorm2d(6),
        nn.Sigmoid(), nn.AvgPool2d(kernel_size=2, stride=2),
        nn.Conv2d(6, 16, kernel_size=5), nn.BatchNorm2d(16),
        nn.Sigmoid(),
        nn.AvgPool2d(kernel_size=2, stride=2),
        nn.Flatten(),
        nn.Linear(16 * 5 * 5, 120), nn.BatchNorm1d(120),
        nn.Sigmoid(),
        nn.Linear(120, 84), nn.BatchNorm1d(84),
        nn.Sigmoid(),
        nn.Linear(84, 10))

    # 初始化参数
    def init_weights(m):
        if type(m) == nn.Linear or type(m) == nn.Conv2d:
            nn.init.xavier_uniform_(m.weight)

    net.apply(init_weights)
    # 模型迁移
    net = net.to(device)
    return net


def train(net, train_iter):
    # optimizer会根据梯度对所有参数进行更新
    optimizer = torch.optim.SGD(net.parameters(), lr)
    loss = nn.CrossEntropyLoss()
    for epoch in range(0, epoch_num):
        net.train()
        total = 0
        correct = 0
        for X, y in train_iter:
            # 梯度置零
            X, y = X.to(device), y.to(device)
            y_hat = net(X)
            l = loss(y_hat, y)
            # 梯度反向传播，参数更新
            optimizer.zero_grad()
            l.backward()
            optimizer.step()
            # 计算准确率
            _, pred = torch.max(y_hat, 1)
            total += y.size(0)
            correct += (pred == y).sum().item()
        print(f'epoch = {epoch}, total = {total}, acurrcy = {100 * correct / total}%')


def test(net, test_iter):
    net.eval()
    with torch.no_grad():
        correct = 0
        total = 0
        for X, y in test_iter:
            X = X.to(device)
            y = y.to(device)
            output = net(X)
            _, pred = torch.max(output, 1)
            total += y.size(0)
            correct += (pred == y).sum().item()
        print(f'test,total = {total}, acurrcy = {100 * correct / total}%')


if __name__ == '__main__':
    train_iter, test_iter = load()
    net = netdefine()
    train(net, train_iter)
    test(net, test_iter)