Pytorch-Learning/train.py at master · xw675/Pytorch-Learning · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
import torch
import torch.nn as nn
import torchvision
from IPython.core.pylabtools import figsize
from torch.utils.data import DataLoader
from CIFAR10_model import MyCIFAR10
from torch.optim import lr_scheduler
from torch.utils.tensorboard import SummaryWriter
import matplotlib.pyplot as plt
import numpy as np
import random

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

mean = (0.4914, 0.4822, 0.4465)
std = (0.247, 0.243, 0.261)
train_transform = torchvision.transforms.Compose([torchvision.transforms.RandomCrop(32, padding=4),
                                                  torchvision.transforms.RandomHorizontalFlip(),
                                                  torchvision.transforms.ToTensor(),
                                                  torchvision.transforms.Normalize(mean, std)])
test_transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor(),
                                                torchvision.transforms.Normalize(mean, std)])

train_dataset = torchvision.datasets.CIFAR10(root='./Dataset', train=True, transform=train_transform, download=False)
test_dataset = torchvision.datasets.CIFAR10(root='./Dataset', train=False, transform=test_transform, download=False)
train_data_loader = DataLoader(train_dataset, batch_size=64)
test_data_loader = DataLoader(test_dataset, batch_size=64)

model = MyCIFAR10().to(device)
model.load_state_dict(torch.load("CIFAR10_model.pth"))
loss = nn.CrossEntropyLoss(label_smoothing=0.1).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-5)
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=3)


writer = SummaryWriter("logs")

def train():
    model.train()
    for epoch in range(30):
        accuracy = 0
        print(f"Epoch {epoch + 1}")
        train_loss = 0.0
        for i , (image, label) in enumerate(train_data_loader):
            image, label = image.to(device), label.to(device)
            images, labels = image, label
            output = model(images)
            result_loss = loss(output, labels)

            optimizer.zero_grad()
            result_loss.backward()
            optimizer.step()
            train_loss += result_loss.item()
            writer.add_scalar("train_loss", result_loss.item(), epoch * len(train_data_loader) + i)
            pred = torch.argmax(output, dim=1)
            accuracy += (pred == labels).sum().item()

        avg_loss = train_loss / len(train_data_loader)
        scheduler.step(avg_loss)

        print(f"Average Loss: {avg_loss}")
        print(f"Accuracy: {accuracy / len(train_dataset)}")

    writer.close()
    torch.save(model.state_dict(), "CIFAR10_model.pth")

def test():
    model.eval()
    accuracy = 0
    with torch.no_grad():
        for i, (image, label) in enumerate(test_data_loader):
            image, label = image.to(device), label.to(device)
            images, labels = image, label
            output = model(images)

            result_loss = loss(output, labels)
            writer.add_scalar("test_loss", result_loss.item(), i * len(test_data_loader) + i)

            pred = torch.argmax(output, dim=1)
            accuracy += (pred == labels).sum().item()

        writer.add_scalar("test_accuracy", accuracy / len(test_dataset), i * len(test_data_loader) + i)
        print(f"Test Accuracy: {accuracy / len(test_dataset)}")

def show():
    class_name = ['Plane', 'Car', 'Bird', 'Cat', 'Deer', 'Dog', 'Frog', 'Horse', 'Ship', 'Truck']
    fig, ax = plt.subplots(2, 6)
    for i in range(2):
        for j in range(6):
            image, label = test_dataset[random.randint(0, len(test_dataset) - 1)]
            image = image.to(device)
            output = model(image.unsqueeze(0))
            pred = torch.argmax(output, dim=1)

            image = image.cpu().numpy().transpose(1, 2, 0)
            image = (image * std + mean) * 255
            image = np.clip(image, 0, 255).astype(np.uint8)

            ax[i][j].imshow(image)
            ax[i][j].axis('off')
            ax[i][j].set_title(f"Pred: {class_name[pred]} \n Label: {class_name[label]}")

    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    train()
    test()
    show()