第一个PyTorch程序训练测试流程

流程说明：

1、数据预处理：transform
2、数据库载入数据：重载Dateset
3、获取训练或者测试的数据流：DateLoader
4、定义网络：重载nn.Module,也有其它更复杂的方法；
5、定义loss：根据自己的问题进行定义；
6、训练流程：定义网络，获取数据流；定义循环，在循环中清空梯度，网络前向传播，计算loss，网络反向传播，梯度更新，并为了监视训练效果，在一定轮数打印出loss的情况；
7、测试流程：定义网络，获取数据流；对每一个数据，传入网络进行网络的前向传播，然后对网络结果进行处理后得到我们想要的结果，计算测试集的loss和准确率，并对相关信息进行打印；



import torch

from torch import nn
from torch import optim
import torch.nn.functional as F
import torchvision


from torch.utils.data import DataLoader, Dataset
import torchvision.transforms as transforms
import cv2
# 数据预处理
def gettransform():
    # img_path = "./ori.jpg"

    train_transformer = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize([5,1]),
        #transforms.RandomResizedCrop(224,scale=(0.5,1.0)),
        #transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        #transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        #transforms.ToTensor()
        #transforms.ToPILImage()
    ])

    ##numpy.ndarray
    # img = cv2.imread(img_path)  # 读取图像
    # img1 = train_transformer(img)
    #
    # img1.show()
    return  train_transformer


# 重载 Dataset 以便获取dataloader，用于训练
class MyDataset(Dataset):
    # TensorDataset继承Dataset, 重载了__init__, __getitem__, __len__
    # 实现将一组Tensor数据对封装成Tensor数据集
    # 能够通过index得到数据集的数据，能够通过len，得到数据集大小

    def __init__(self, imagepath,mytransformer):
        self.imagepath = imagepath
        self.transformer = mytransformer
        self.testy = 1

    def __getitem__(self,index):
        img = cv2.imread(self.imagepath)  # 读取图像
        img1 = self.transformer(img)
        return img1

    def __len__(self):
        return len(self.imagepath)

train_batch_size = 64
test_batch_size = 1000
img_size = 28


def get_dataloader(train=True):
    assert isinstance(train, bool), "train 必须是bool类型"

    # 准备数据集，其中0.1307，0.3081为MNIST数据的均值和标准差，这样操作能够对其进行标准化
    # 因为MNIST只有一个通道（黑白图片）,所以元组中只有一个值
    dataset = torchvision.datasets.MNIST('/data', train=train, download=True,
                                         transform=torchvision.transforms.Compose([
                                             torchvision.transforms.ToTensor(),
                                             torchvision.transforms.Normalize((0.1307,), (0.3081,)), ]))
    # 准备数据迭代器
    batch_size = train_batch_size if train else test_batch_size
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
    return dataloader


class MnistNet(nn.Module):
    def __init__(self):
        super(MnistNet, self).__init__()
        self.fc1 = nn.Linear(28 * 28 * 1, 28)
        self.fc2 = nn.Linear(28, 10)

    def forward(self, x):
        x = x.view(-1, 28 * 28 * 1)
        x = self.fc1(x)  # [batch_size,28]
        x = F.relu(x)  # [batch_size,28]
        x = self.fc2(x)  # [batch_size,10]
        # return x
        return F.log_softmax(x, dim=-1)


mnist_net = MnistNet()
optimizer = optim.Adam(mnist_net.parameters(), lr=0.001)

train_loss_list = []
train_count_list = []

def train(epoch):
    mode = True
    mnist_net.train(mode=mode)
    train_dataloader = get_dataloader(train=mode)
    print(len(train_dataloader.dataset))
    print(len(train_dataloader))
    for idx, (data, target) in enumerate(train_dataloader):
        optimizer.zero_grad()
        output = mnist_net(data)
        loss = F.nll_loss(output, target)  # 对数似然损失
        loss.backward()
        optimizer.step()
        if idx % 10 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, idx * len(data), len(train_dataloader.dataset),
                       100. * idx / len(train_dataloader), loss.item()))

            train_loss_list.append(loss.item())
            train_count_list.append(idx * train_batch_size + (epoch - 1) * len(train_dataloader))
            torch.save(mnist_net.state_dict(), "model/mnist_net.pkl")
            torch.save(optimizer.state_dict(), 'results/mnist_optimizer.pkl')


def test():
    test_loss = 0
    correct = 0
    mnist_net.eval()
    test_dataloader = get_dataloader(train=False)
    with torch.no_grad():
        for data, target in test_dataloader:
            output = mnist_net(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()
            pred = output.data.max(1, keepdim=True)[1]  # 获取最大值的位置,[batch_size,1]
            correct += pred.eq(target.data.view_as(pred)).sum()
    test_loss /= len(test_dataloader.dataset)
    print('\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n'.format(
        test_loss, correct, len(test_dataloader.dataset),
        100. * correct / len(test_dataloader.dataset)))


# train and test
if __name__ == '__main__':

    test()
    for i in range(5):  # 模型训练5轮
        train(i)
        test()