Vision Transformer(ViT)实践项目，图像分类任务，“猫狗大战”（猫狗分类）

from __future__ import print_function import glob from itertools import chain import os import random import zipfile import matplotlib.pyplot as plt import numpy as np import pandas as pd import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from linformer import Linformer from PIL import Image from sklearn.model_selection import train_test_split from torch.optim.lr_scheduler import StepLR from torch.utils.data import DataLoader, Dataset from torchvision import datasets, transforms from tqdm.notebook import tqdm from vit_pytorch.vit import ViT # Training settings batch_size = 64 epochs = 10 lr = 3e-5 gamma = 0.7 seed = 42 print(torch.version.cuda) print(torch.__version__) print("cuda:" , torch.cuda.is_available()) def seed_everything(seed): random.seed(seed) os.environ['PYTHONHASHSEED'] = str(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True seed_everything(seed) device = 'cuda' os.makedirs('data', exist_ok=True) train_dir = './data/train' test_dir = './data/test' train_list1 = glob.glob(os.path.join(train_dir+'/cats', '*.jpg')) train_list2 = glob.glob(os.path.join(train_dir+'/dogs', '*.jpg')) train_list = train_list1+train_list2 test_list1 = glob.glob(os.path.join(test_dir+'/cats', '*.jpg')) test_list2 = glob.glob(os.path.join(test_dir+'/dogs', '*.jpg')) test_list = test_list1+test_list2 print(f"Train Data: {len(train_list)}") print(f"Test Data: {len(test_list)}") labels = [path.split('/')[-1].split('.')[0] for path in train_list] train_list, valid_list = train_test_split(train_list, test_size=0.2, stratify=labels, random_state=seed) print(f"Train Data: {len(train_list)}") print(f"Validation Data: {len(valid_list)}") print(f"Test Data: {len(test_list)}") train_transforms = transforms.Compose( [ transforms.Resize((224, 224)), transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), ] ) val_transforms = transforms.Compose( [ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), ] ) test_transforms = transforms.Compose( [ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), ] ) class CatsDogsDataset(Dataset): def __init__(self, file_list, transform=None): self.file_list = file_list self.transform = transform def __len__(self): self.filelength = len(self.file_list) return self.filelength def __getitem__(self, idx): img_path = self.file_list[idx] img = Image.open(img_path) img_transformed = self.transform(img) label = img_path.split("/")[-1].split(".")[0] label = 1 if label == "dog" else 0 return img_transformed, label train_data = CatsDogsDataset(train_list, transform=train_transforms) valid_data = CatsDogsDataset(valid_list, transform=test_transforms) test_data = CatsDogsDataset(test_list, transform=test_transforms) train_loader = DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True) valid_loader = DataLoader(dataset=valid_data, batch_size=batch_size, shuffle=True) test_loader = DataLoader(dataset=test_data, batch_size=batch_size, shuffle=True) print(len(train_data), len(train_loader)) print(len(valid_data), len(valid_loader)) model = ViT( image_size = 256, patch_size = 16, num_classes =2, dim = 1024, depth = 6, heads = 16, mlp_dim = 2048, dropout = 0.1, emb_dropout = 0.1 ).to(device) print(model) # from torchsummary import summary # from ptflops import get_model_complexity_info # # summary(model,(3,256,256)) # get_model_complexity_info(model,(3,256,256)) # img = torch.randn(1,3, 256, 256) # preds = model(img) # print(preds.shape) # loss function criterion = nn.CrossEntropyLoss() # optimizer optimizer = optim.Adam(model.parameters(), lr=lr) # scheduler scheduler = StepLR(optimizer, step_size=1, gamma=gamma) best_accuracy = 0 if os.path.exists('best.mdl'): model.load_state_dict(torch.load('best.mdl')) for epoch in range(epochs): #train epoch_loss = 0 epoch_accuracy = 0 for data, label in tqdm(train_loader): data = data.to(device) label = label.to(device) output = model(data) loss = criterion(output, label) optimizer.zero_grad() loss.backward() optimizer.step() acc = (output.argmax(dim=1) == label).float().mean() epoch_accuracy += acc / len(train_loader) epoch_loss += loss / len(train_loader) #validataion with torch.no_grad(): epoch_val_accuracy = 0 epoch_val_loss = 0 for data, label in valid_loader: data = data.to(device) label = label.to(device) val_output = model(data) val_loss = criterion(val_output, label) acc = (val_output.argmax(dim=1) == label).float().mean() epoch_val_accuracy += acc / len(valid_loader) epoch_val_loss += val_loss / len(valid_loader) if epoch_val_accuracy>best_accuracy: best_accuracy = epoch_val_accuracy torch.save(model.state_dict(), 'best.mdl') print( f"Epoch : {epoch + 1} - loss : {epoch_loss:.4f} - acc: {epoch_accuracy:.4f} - val_loss : {epoch_val_loss:.4f} - val_acc: {epoch_val_accuracy:.4f}\n" ) #test model.load_state_dict(torch.load('best.mdl')) print('best_acc:',best_accuracy) with torch.no_grad(): test_accuracy = 0 test_loss = 0 for data, label in test_loader: data = data.to(device) label = label.to(device) test_output = model(data) test_loss = criterion(test_output, label) acc = (test_output.argmax(dim=1) == label).float().mean() test_accuracy += acc / len(test_loader) test_loss += test_loss / len(test_loader) print('test_acc:',test_accuracy)