实战手写数字识别—使用visdom首先loss函数值可视化

墨蓝 2023-06-07 05:54 30阅读 0赞

首先使用 python -m visdom.server激活visdom环境，再运行以下代码。

import torch
    import torch.nn as nn
    import torchvision.datasets as normal_datasets
    import torchvision.transforms as transforms
    from torch.autograd import Variable
    import visdom
    import numpy as np
    import matplotlib.pyplot as plt
    
    num_epochs = 5
    batch_size = 100
    learning_rate = 0.001
    
    
    # 将数据处理成Variable, 如果有GPU, 可以转成cuda形式
    def get_variable(x):
        x = Variable(x)
        return x.cuda() if torch.cuda.is_available() else x
    
    
    # 从torchvision.datasets中加载一些常用数据集
    train_dataset = normal_datasets.MNIST(
        root='./data',  # 数据集保存路径
        train=True,  # 是否作为训练集
        transform=transforms.ToTensor(),  # 数据如何处理, 可以自己自定义
        download=True)  # 路径下没有的话, 可以下载
    
    # 见数据加载器和batch
    test_dataset = normal_datasets.MNIST(root='./data',
                                         train=False,
                                         transform=transforms.ToTensor(),
                                         download=True)
    
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                               batch_size=batch_size,
                                               shuffle=True)
    
    test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                              batch_size=batch_size,
                                              shuffle=False)
    
    viz = visdom.Visdom(env='train-mnist')
    loss_win = viz.line(np.arange(10))
    # 两层卷积
    class CNN(nn.Module):
        def __init__(self):
            super(CNN, self).__init__()
            # 使用序列工具快速构建
            self.conv1 = nn.Sequential(
                nn.Conv2d(1, 16, kernel_size=5, padding=2),
                nn.BatchNorm2d(16),
                nn.ReLU(),
                nn.MaxPool2d(2))
            self.conv2 = nn.Sequential(
                nn.Conv2d(16, 32, kernel_size=5, padding=2),
                nn.BatchNorm2d(32),
                nn.ReLU(),
                nn.MaxPool2d(2))
            self.fc = nn.Linear(7 * 7 * 32, 10)
    
        def forward(self, x):
            out = self.conv1(x)
            out = self.conv2(out)
            out = out.view(out.size(0), -1)  # reshape
            out = self.fc(out)
            return out
    
    
    cnn = CNN()
    if torch.cuda.is_available():
        cnn = cnn.cuda()
    
    # 选择损失函数和优化方法
    loss_func = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(cnn.parameters(), lr=learning_rate)
    
    loss_point = []  ####
    x_point = []   ####
    iter_count = 0  ####
    
    for epoch in range(num_epochs):
    
        running_loss = 0.0  ####
        for i, (images, labels) in enumerate(train_loader):
    
    
            images = get_variable(images)
            labels = get_variable(labels)
    
            outputs = cnn(images)
            loss = loss_func(outputs, labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    
            running_loss += loss.item()  ####
    
            if (i + 1) % 100 == 0:
                print('Epoch [%d/%d], Iter [%d/%d] Loss: %.4f'
                      % (epoch + 1, num_epochs, i + 1, len(train_dataset) // batch_size, loss.item()))
    
                tr_loss = running_loss / 100
                iter_count += 100
                if iter_count == 100:
                    viz.line(Y=np.array([tr_loss]), X=np.array([iter_count]), update='replace', win=loss_win)
                else:
                    viz.line(Y=np.array([tr_loss]), X=np.array([iter_count]), update='append', win=loss_win)
            running_loss = 0 ####
    
    
    # Save the Trained Model
    torch.save(cnn.state_dict(), 'cnn.pkl')