在CIFAR-10上训练VGG6

import os
import torch
from torch.utils.data import Dataset
from torchvision.io import read_image
from torch.utils.data import DataLoader
import sys
import torchvision.transforms as transforms
import torch.optim as optim
import torch.nn as nn
from tqdm import tqdm
from typing import List, cast

classes = ('airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
class Cinic10(Dataset):
    def __init__(self, img_dir):
        self.img_labels = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
        self.label2id = {'airplane': 0, 'automobile': 1, 'bird': 2, 'cat': 3, 'deer': 4,
                         'dog': 5, 'frog': 6, 'horse': 7, 'ship': 8, 'truck': 9}
        self.id2label = {0: 'airplane', 1: 'automobile', 2: 'bird', 3: 'cat', 4: 'deer',
                         5: 'dog', 6: 'frog', 7: 'horse', 8: 'ship', 9: 'truck'}
        self.img_dir = img_dir
        self.cinic_mean = [0.47889522, 0.47227842, 0.43047404]
        self.cinic_std = [0.24205776, 0.23828046, 0.25874835]
        self.transform = transforms.Compose([transforms.Normalize(mean=self.cinic_mean, std=self.cinic_std)])
        self.countPerLabel = len(
            [name for name in os.listdir(os.path.join(self.img_dir, self.img_labels[0]))
             if os.path.isfile(os.path.join(os.path.join(self.img_dir, self.img_labels[0]), name))])
        self.len = len(self.img_labels) * self.countPerLabel
        self.X_Y = []
        for label in self.img_labels:
            img_path = os.path.join(self.img_dir, label)
            images_files = [name for name in os.listdir(img_path) if os.path.isfile(os.path.join(img_path, name))]
            label_id = self.label2id[label]
            for images_file in images_files:
                image = read_image(os.path.join(img_path, images_file))
                if image.shape != torch.Size([3, 32, 32]):
                    image = torch.cat([image, image, image])
                image = image.type(torch.float32)
                image = self.transform(image)
                self.X_Y.append([image, label_id])

    def __len__(self):
        return self.len

    def __getitem__(self, idx):
        [image, label] = self.X_Y[idx]
        return image, label


class Cifar10(Dataset):
    def __init__(self, img_dir):
        self.img_labels = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
        self.label2id = {'airplane': 0, 'automobile': 1, 'bird': 2, 'cat': 3, 'deer': 4,
                         'dog': 5, 'frog': 6, 'horse': 7, 'ship': 8, 'truck': 9}
        self.id2label = {0: 'airplane', 1: 'automobile', 2: 'bird', 3: 'cat', 4: 'deer',
                         5: 'dog', 6: 'frog', 7: 'horse', 8: 'ship', 9: 'truck'}
        self.img_dir = img_dir
        self.cinic_mean = [0.485, 0.456, 0.406]
        self.cinic_std = [0.229, 0.224, 0.225]
        self.transform = transforms.Compose([transforms.Normalize(mean=self.cinic_mean, std=self.cinic_std)])
        self.countPerLabel = len(
            [name for name in os.listdir(os.path.join(self.img_dir, self.img_labels[0]))
             if os.path.isfile(os.path.join(os.path.join(self.img_dir, self.img_labels[0]), name))])
        self.len = len(self.img_labels) * self.countPerLabel
        self.X_Y = []
        for label in self.img_labels:
            img_path = os.path.join(self.img_dir, label)
            images_files = [name for name in os.listdir(img_path) if os.path.isfile(os.path.join(img_path, name))]
            label_id = self.label2id[label]
            for images_file in images_files:
                image = read_image(os.path.join(img_path, images_file))
                if image.shape != torch.Size([3, 32, 32]):
                    image = torch.cat([image, image, image])
                image = image.type(torch.float32)
                image = self.transform(image)
                self.X_Y.append([image, label_id])

    def __len__(self):
        return self.len

    def __getitem__(self, idx):
        [image, label] = self.X_Y[idx]
        return image, label


traindataset = Cifar10('CIFAR10/train')
traindataloader = DataLoader(traindataset, batch_size=256, shuffle=True)
testdataset = Cifar10('CIFAR10/test')
testdataloader = DataLoader(testdataset, batch_size=256)


def make_layers():
    layers: List[nn.Module] = []
    in_channels = 3
    cfg = [64, 64, "M", 128, 128, "M", 256, 256, 256, "M", 512, 512, 512, "M", 512, 512, 512, "M"]
    for v in cfg:
        if v == "M":
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            v = cast(int, v)
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            in_channels = v
    return nn.Sequential(*layers)


class VGG(nn.Module):
    def __init__(self, features, num_classes, dropout):
        super().__init__()
        self.features = features
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(True),
            nn.Dropout(p=dropout),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(p=dropout),
            nn.Linear(4096, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x


VGG16_layers = [64, 64, "M", 128, 128, "M", 256, 256, 256, "M", 512, 512, 512, "M", 512, 512, 512, "M"]
device = torch.device('cuda:1')
net = VGG(make_layers(), len(classes), 0.5)
net = net.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.05, momentum=0.9, weight_decay=5e-4)


def adjust_learning_rate(optimizer, epoch):
    """Sets the learning rate to the initial LR decayed by 2 every 30 epochs"""
    lr = 0.05 * (0.5 ** (epoch // 30))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr


def train(dataloader, net, criterion, optimizer, device):
    net.train()
    running_loss = 0.0
    correct = 0
    total = 0
    for data in tqdm(dataloader, desc='training...', file=sys.stdout):
        inputs, labels = data
        inputs = inputs.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        running_loss += loss.item()
    return running_loss/total, correct/total

def evaluate(dataloader, net, criterion, device):
    net.eval()
    running_loss = 0.0
    correct = 0
    total = 0
    with torch.no_grad():
        for data in tqdm(dataloader, desc='evaluating...', file=sys.stdout):
            inputs, labels = data
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            running_loss += loss.item()
    return running_loss/total, correct/total


n_epochs = 300
best_valid_acc = 0
for epoch in range(n_epochs):
    adjust_learning_rate(optimizer, epoch)
    train_loss, train_acc = train(traindataloader, net, criterion, optimizer, device)
    valid_loss, valid_acc = evaluate(testdataloader, net, criterion, device)
    print(f'epoch: {epoch + 1}')
    print(f'train_loss: {train_loss:.3f}, train_acc: {train_acc:.3f}')
    print(f'valid_loss: {valid_loss:.3f}, valid_acc: {valid_acc:.3f}')
    if valid_acc > best_valid_acc:
        print(f'{valid_acc:.3f} is better than {best_valid_acc:.3f}, best valid acc is {valid_acc:.3f}')
        best_valid_acc = valid_acc
        torch.save(net.state_dict(), 'CV/CIFAR10/VGG16.pth')
    else:
        print(f'best valid acc is {best_valid_acc:.3f}')

net2 = VGG(make_layers(), len(classes), 0.5)
net2 = net2.to(device)
net2.load_state_dict(torch.load('CV/CIFAR10/VGG16.pth'))
valid_loss, valid_acc = evaluate(testdataloader, net2, criterion, device)
print(f'best model valid acc: {valid_acc:.3f}')