PoolFormer实战：使用PoolFormer实现图像分类任务.zip

import json import os import matplotlib.pyplot as plt import torch import torch.nn as nn import torch.nn.parallel import torch.optim as optim import torch.utils.data import torch.utils.data.distributed import torchvision.transforms as transforms from timm.utils import accuracy, AverageMeter, ModelEma from sklearn.metrics import classification_report from timm.data.mixup import Mixup from timm.loss import SoftTargetCrossEntropy from timm.models import poolformer_s24 from torch.autograd import Variable from torchvision import datasets torch.backends.cudnn.benchmark = False import warnings warnings.filterwarnings("ignore") os.environ['CUDA_VISIBLE_DEVICES']="0,1" # 定义训练过程 def train(model, device, train_loader, optimizer, epoch,model_ema): model.train() loss_meter = AverageMeter() acc1_meter = AverageMeter() acc5_meter = AverageMeter() total_num = len(train_loader.dataset) print(total_num, len(train_loader)) for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device, non_blocking=True), Variable(target).to(device, non_blocking=True) samples, targets = mixup_fn(data, target) output = model(samples) optimizer.zero_grad() if use_amp: with torch.cuda.amp.autocast(): loss = torch.nan_to_num(criterion_train(output, targets)) scaler.scale(loss).backward() torch.nn.utils.clip_grad_norm_(model.parameters(), CLIP_GRAD) # Unscales gradients and calls # or skips optimizer.step() scaler.step(optimizer) # Updates the scale for next iteration scaler.update() else: loss = criterion_train(output, targets) loss.backward() # torch.nn.utils.clip_grad_norm_(model.parameters(), CLIP_GRAD) optimizer.step() if model_ema is not None: model_ema.update(model) torch.cuda.synchronize() lr = optimizer.state_dict()['param_groups'][0]['lr'] loss_meter.update(loss.item(), target.size(0)) acc1, acc5 = accuracy(output, target, topk=(1, 5)) loss_meter.update(loss.item(), target.size(0)) acc1_meter.update(acc1.item(), target.size(0)) acc5_meter.update(acc5.item(), target.size(0)) if (batch_idx + 1) % 10 == 0: print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tLR:{:.9f}'.format( epoch, (batch_idx + 1) * len(data), len(train_loader.dataset), 100. * (batch_idx + 1) / len(train_loader), loss.item(), lr)) ave_loss =loss_meter.avg acc = acc1_meter.avg print('epoch:{}\tloss:{:.2f}\tacc:{:.2f}'.format(epoch, ave_loss, acc)) return ave_loss, acc # 验证过程 @torch.no_grad() def val(model, device, test_loader): global Best_ACC model.eval() loss_meter = AverageMeter() acc1_meter = AverageMeter() acc5_meter = AverageMeter() total_num = len(test_loader.dataset) print(total_num, len(test_loader)) val_list = [] pred_list = [] for data, target in test_loader: for t in target: val_list.append(t.data.item()) data, target = data.to(device,non_blocking=True), target.to(device,non_blocking=True) output = model(data) loss = criterion_val(output, target) _, pred = torch.max(output.data, 1) for p in pred: pred_list.append(p.data.item()) acc1, acc5 = accuracy(output, target, topk=(1, 5)) loss_meter.update(loss.item(), target.size(0)) acc1_meter.update(acc1.item(), target.size(0)) acc5_meter.update(acc5.item(), target.size(0)) acc = acc1_meter.avg print('\nVal set: Average loss: {:.4f}\tAcc1:{:.3f}%\tAcc5:{:.3f}%\n'.format( loss_meter.avg, acc, acc5_meter.avg)) if acc > Best_ACC: if isinstance(model, torch.nn.DataParallel): torch.save(model.module, file_dir + '/' + 'best.pth') else: torch.save(model, file_dir + '/' + 'best.pth') Best_ACC = acc if isinstance(model, torch.nn.DataParallel): state = { 'epoch': epoch, 'state_dict': model.module.state_dict(), 'Best_ACC':Best_ACC } if use_ema: state['state_dict_ema']=model.module.state_dict() torch.save(state, file_dir + "/" + 'model_' + str(epoch) + '_' + str(round(acc, 3)) + '.pth') else: state = { 'epoch': epoch, 'state_dict': model.state_dict(), 'Best_ACC': Best_ACC } if use_ema: state['state_dict_ema']=model.state_dict() torch.save(state, file_dir + "/" + 'model_' + str(epoch) + '_' + str(round(acc, 3)) + '.pth') return val_list, pred_list, loss_meter.avg, acc if __name__ == '__main__': #创建保存模型的文件夹 file_dir = 'checkpoints/PoolFormer' if os.path.exists(file_dir): print('true') os.makedirs(file_dir,exist_ok=True) else: os.makedirs(file_dir) # 设置全局参数 model_lr = 1e-4 BATCH_SIZE = 16 EPOCHS = 1000 DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') use_amp = True # 是否使用混合精度 use_dp = True #是否开启dp方式的多卡训练 classes = 12 resume =None CLIP_GRAD = 5.0 Best_ACC = 0 #记录最高得分 use_ema=True model_ema_decay=0.9998 start_epoch=1 # 数据预处理7 transform = transforms.Compose([ transforms.RandomRotation(10), transforms.GaussianBlur(kernel_size=(5,5),sigma=(0.1, 3.0)), transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5), transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.44127703, 0.4712498, 0.43714803], std= [0.18507297, 0.18050247, 0.16784933]) ]) transform_test = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.44127703, 0.4712498, 0.43714803], std= [0.18507297, 0.18050247, 0.16784933]) ]) mixup_fn = Mixup( mixup_alpha=0.8, cutmix_alpha=1.0, cutmix_minmax=None, prob=0.1, switch_prob=0.5, mode='batch', label_smoothing=0.1, num_classes=classes) # 读取数据 dataset_train = datasets.ImageFolder('data/train', transform=transform) dataset_test = datasets.ImageFolder("data/val", transform=transform_test) with open('class.txt', 'w') as file: file.write(str(dataset_train.class_to_idx)) with open('class.json', 'w', encoding='utf-8') as file: file.write(json.dumps(dataset_train.class_to_idx)) # 导入数据 train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True,drop_last=True) test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=False) # 实例化模型并且移动到GPU criterion_train = SoftTargetCrossEntropy() criterion_val = torch.nn.CrossEntropyLoss() #设置模型 model_ft = poolformer_s24(pretrained=True) model_ft.reset_classifier(classes) if resume: model=torch.load(resume) print(model['state_dict'].keys()) model_ft.load_state_dict(model['state_dict']) Best_ACC=model['Best_ACC'] start_epoch=model['epoch']+1 model_ft.to(DEVICE) # 选择简单暴力的Adam优化器，学习率调低 optimizer = optim.AdamW(model_ft.parameters(),lr=model_lr) cosine_schedule = optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer, T_max=20, eta_min=1e-6) if use_amp: scaler = torch.cuda.amp.GradScaler() if torch.cuda.device_count() > 1 and use_dp: print("Let's use", torch.cuda.device_count(), "GPUs!") model_ft = torch.nn.DataParallel(model_ft) if use_ema