DEiT实战:使用DEiT实现图像分类任务(二)
文章目录
- 训练
- 运行以及结果查看
- 测试
- 完整的代码
在上一篇文章中完成了前期的准备工作,见链接:
DEiT实战:使用DEiT实现图像分类任务(一)
这篇主要是讲解如何训练和测试
训练
完成上面的步骤后,就开始train脚本的编写,新建train.py和train_dist.py
导入项目使用的库
在train.py导入
import jsonimport osimport shutilimport matplotlib.pyplot as pltimport torchimport torch.nn as nnimport torch.nn.parallelimport torch.optim as optimimport torch.utils.dataimport torch.utils.data.distributedimport torchvision.transforms as transformsfrom timm.utils import accuracy, AverageMeterfrom sklearn.metrics import classification_reportfrom timm.data.mixup import Mixupfrom timm.loss import SoftTargetCrossEntropyfrom torchvision import datasetsfrom timm.models import deit_small_patch16_224torch.backends.cudnn.benchmark = Falseimport warningswarnings.filterwarnings("ignore")from ema import EMA
在train_dist.py导入
import jsonimport osimport matplotlib.pyplot as pltimport torchimport torch.nn as nnimport torch.nn.parallelimport torch.optim as optimimport torch.utils.dataimport torch.utils.data.distributedimport torchvision.transforms as transformsfrom timm.utils import accuracy, AverageMeterfrom sklearn.metrics import classification_reportfrom timm.data.mixup import Mixupfrom timm.loss import SoftTargetCrossEntropyfrom torchvision import datasetsfrom timm.models import deit_small_distilled_patch16_224torch.backends.cudnn.benchmark = Falseimport warningswarnings.filterwarnings("ignore")from ema import EMA
distilled表示含有蒸馏的token。
设置随机因子
def seed_everything(seed=42): os.environ['PYHTONHASHSEED'] = str(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) torch.backends.cudnn.deterministic = True
设置了固定的随机因子,再次训练的时候就可以保证图片的加载顺序不会发生变化。
设置全局参数
设置学习率、BatchSize、epoch等参数,判断环境中是否存在GPU,如果没有则使用CPU。建议使用GPU,CPU太慢了。
if __name__ == '__main__': #创建保存模型的文件夹 file_dir = 'checkpoints/DEiT' 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=False #是否开启dp方式的多卡训练 classes = 12 resume = False CLIP_GRAD = 5.0 model_path = 'best.pth' Best_ACC = 0 #记录最高得分 use_ema=True SEED=42 seed_everything(42)
设置存放权重文件的文件夹,如果文件夹存在删除再建立。
接下来,查看全局参数:
model_lr:学习率,根据实际情况做调整。
BATCH_SIZE:batchsize,根据显卡的大小设置。
EPOCHS:epoch的个数,一般300够用。
use_amp:是否使用混合精度。
classes:类别个数。
resume:是否接着上次模型继续训练。
model_path:模型的路径。如果resume设置为True时,就采用model_path定义的模型继续训练。
CLIP_GRAD:梯度的最大范数,在梯度裁剪里设置。
Best_ACC:记录最高ACC得分。
use_ema:是否使用ema
SEED:随机因子,数值可以随意设定,但是设置后,不要随意更改,更改后,图片加载的顺序会改变,影响测试结果。
file_dir = 'checkpoints/DEiT'
这是存放DEiT模型的路径。
在train_dist.py 则应该设置为:
file_dir = 'checkpoints/DEiT_dist'
图像预处理与增强
数据处理比较简单,加入了Cutout、做了Resize和归一化,定义Mixup函数。
这里注意下Resize的大小,由于选用的MaxViT模型输入是224×224的大小,所以要Resize为224×224。
# 数据预处理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)
读取数据
使用pytorch默认读取数据的方式,然后将dataset_train.class_to_idx打印出来,预测的时候要用到。
将dataset_train.class_to_idx保存到txt文件或者json文件中。
# 读取数据 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) test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=False)
class_to_idx的结果:
{‘Black-grass’: 0, ‘Charlock’: 1, ‘Cleavers’: 2, ‘Common Chickweed’: 3, ‘Common wheat’: 4, ‘Fat Hen’: 5, ‘Loose Silky-bent’: 6, ‘Maize’: 7, ‘Scentless Mayweed’: 8, ‘Shepherds Purse’: 9, ‘Small-flowered Cranesbill’: 10, ‘Sugar beet’: 11}
设置模型
train.py
- 设置loss函数,train的loss为:SoftTargetCrossEntropy,val的loss:nn.CrossEntropyLoss()。
- 设置模型为deit_small_patch16_224,pretrained设置为true,表示加载预训练模型,调用reset_classifier函数将classes设置为12。如果resume为True,则加载模型接着上次训练。
- 优化器设置为adamW。
- 学习率调整策略选择为余弦退火。
- 开启混合精度训练,声明pytorch自带的混合精度 torch.cuda.amp.GradScaler()。
- 检测可用显卡的数量,如果大于1,并且开启多卡训练的情况下,则要用torch.nn.DataParallel加载模型,开启多卡训练。
- 如果使用ema,则注册ema
# 实例化模型并且移动到GPU criterion_train = SoftTargetCrossEntropy() criterion_val = torch.nn.CrossEntropyLoss() #设置模型 model_ft = deit_small_patch16_224(pretrained=True) model_ft.reset_classifier(classes) # num_ftrs = model_ft.head.in_features # model_ft.head = nn.Linear(num_ftrs, classes) if resume: model = torch.load(resume) model_ft.load_state_dict(model['state_dict']) Best_ACC = model['Best_ACC'] start_epoch = model['epoch'] + 1 model_ft.to(DEVICE) print(model_ft) # 选择简单暴力的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: ema = EMA(model_ft, 0.9998) ema.register()
注:torch.nn.DataParallel方式,默认不能开启混合精度训练的,如果想要开启混合精度训练,则需要在模型的forward前面加上@autocast()函数。
如果不开启混合精度则要将@autocast()去掉,否则loss一直试nan。
train_dist.py
train_dist.py设置模型为deit_small_distilled_patch16_224。
#设置模型 model_ft = deit_small_distilled_patch16_224(pretrained=True) model_ft.reset_classifier(classes)
定义训练和验证函数
训练函数
训练的主要步骤:
1、使用AverageMeter保存自定义变量,包括loss,ACC1,ACC5。
2、判断迭代的数据是否是奇数,由于mixup_fn只能接受偶数,所以如果不是偶数则要减去一位,让其变成偶数。但是有可能最后一次迭代只有一条数据,减去后就变成了0,所以还要判断不能小于2,如果小于2则直接中断本次循环。
3、将数据输入mixup_fn生成mixup数据,然后输入model计算loss。
4、 optimizer.zero_grad() 梯度清零,把loss关于weight的导数变成0。
5、如果使用混合精度,则
- with torch.cuda.amp.autocast(),开启混合精度。
- 计算loss。
- scaler.scale(loss).backward(),梯度放大。
- torch.nn.utils.clip_grad_norm_,梯度裁剪,放置梯度爆炸。
- scaler.step(optimizer) ,首先把梯度值unscale回来,如果梯度值不是inf或NaN,则调用optimizer.step()来更新权重,否则,忽略step调用,从而保证权重不更新。
- 更新下一次迭代的scaler。
否则,直接反向传播求梯度。torch.nn.utils.clip_grad_norm_函数执行梯度裁剪,防止梯度爆炸。
6、如果use_ema为True,则执行model_ema的updata函数,更新模型。
7、 torch.cuda.synchronize(),等待上面所有的操作执行完成。
8、接下来,更新loss,ACC1,ACC5的值。
等待一个epoch训练完成后,计算平均loss和平均acc
# 定义训练过程def train(model, device, train_loader, optimizer, epoch): 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), 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 = 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 use_ema : ema.update() 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
验证函数
验证集和训练集大致相似,主要步骤:
1、定义参数,loss_meter 测试的loss,total_num总的验证集的数量,val_list验证集的label,pred_list预测的label。
2、在val的函数上面添加@torch.no_grad(),作用:所有计算得出的tensor的requires_grad都自动设置为False。即使一个tensor(命名为x)的requires_grad = True,在with torch.no_grad计算,由x得到的新tensor(命名为w-标量)requires_grad也为False,且grad_fn也为None,即不会对w求导。
3、如果use_ema 为True,则使用shadow字典的参数更新模型参数。
4、使用验证集的loss函数求出验证集的loss。
5、调用accuracy函数计算ACC1和ACC56、更新loss_meter、acc1_meter、acc5_meter的参数。
7、如果use_ema 为True,则清空backup字典。
本次epoch循环完成后,求得本次epoch的acc、loss。
如果acc比Best_ACC大,则保存模型。
保存模型的逻辑:
如果ACC比Best_ACC高,则保存best模型
判断模型是否为DP方式训练的模型。如果是DP方式训练的模型,模型参数放在model.module,则需要保存model.module。
否则直接保存model。接下来保存每个epoch的模型。
判断模型是否为DP方式训练的模型。如果是DP方式训练的模型,模型参数放在model.module,则需要保存model.module.state_dict()。
新建个字典,放置Best_ACC、epoch和 model.module.state_dict()等参数。然后将这个字典保存。
否则,新建个字典,放置Best_ACC、epoch和 model.state_dict()等参数。然后将这个字典保存。
在这里注意:对于每个epoch的模型只保存了state_dict参数,没有保存整个模型文件。
# 验证过程@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 = [] if use_ema : ema.apply_shadow() 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)) if use_ema : ema.restore() 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 } 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 } torch.save(state, file_dir + "/" + 'model_' + str(epoch) + '_' + str(round(acc, 3)) + '.pth') return val_list, pred_list, loss_meter.avg, acc
调用训练和验证方法
调用训练函数和验证函数的主要步骤:
1、定义参数:
- is_set_lr,是否已经设置了学习率,当epoch大于一定的次数后,会将学习率设置到一定的值,并将其置为True。
- log_dir:记录log用的,将有用的信息保存到字典中,然后转为json保存起来。
- train_loss_list:保存每个epoch的训练loss。
- val_loss_list:保存每个epoch的验证loss。
- train_acc_list:保存每个epoch的训练acc。
- val_acc_list:保存么每个epoch的验证acc。
- epoch_list:存放每个epoch的值。
循环epoch
1、调用train函数,得到 train_loss, train_acc,并将分别放入train_loss_list,train_acc_list,然后存入到logdir字典中。
2、调用验证函数,得到val_list, pred_list, val_loss, val_acc。将val_loss, val_acc分别放入val_loss_list和val_acc_list中,然后存入到logdir字典中。
3、保存log。
4、打印本次的测试报告。
5、如果epoch大于600,将学习率设置为固定的1e-6。
6、绘制loss曲线和acc曲线。
# 训练与验证 is_set_lr = False log_dir = {} train_loss_list, val_loss_list, train_acc_list, val_acc_list, epoch_list = [], [], [], [], [] for epoch in range(1, EPOCHS + 1): epoch_list.append(epoch) train_loss, train_acc = train(model_ft, DEVICE, train_loader, optimizer, epoch) train_loss_list.append(train_loss) train_acc_list.append(train_acc) log_dir['train_acc'] = train_acc_list log_dir['train_loss'] = train_loss_list val_list, pred_list, val_loss, val_acc = val(model_ft, DEVICE, test_loader) val_loss_list.append(val_loss) val_acc_list.append(val_acc) log_dir['val_acc'] = val_acc_list log_dir['val_loss'] = val_loss_list log_dir['best_acc'] = Best_ACC with open(file_dir + '/result.json', 'w', encoding='utf-8') as file: file.write(json.dumps(log_dir)) print(classification_report(val_list, pred_list, target_names=dataset_train.class_to_idx)) if epoch < 600: cosine_schedule.step() else: if not is_set_lr: for param_group in optimizer.param_groups: param_group["lr"] = 1e-6 is_set_lr = True fig = plt.figure(1) plt.plot(epoch_list, train_loss_list, 'r-', label=u'Train Loss') # 显示图例 plt.plot(epoch_list, val_loss_list, 'b-', label=u'Val Loss') plt.legend(["Train Loss", "Val Loss"], loc="upper right") plt.xlabel(u'epoch') plt.ylabel(u'loss') plt.title('Model Loss ') plt.savefig(file_dir + "/loss.png") plt.close(1) fig2 = plt.figure(2) plt.plot(epoch_list, train_acc_list, 'r-', label=u'Train Acc') plt.plot(epoch_list, val_acc_list, 'b-', label=u'Val Acc') plt.legend(["Train Acc", "Val Acc"], loc="lower right") plt.title("Model Acc") plt.ylabel("acc") plt.xlabel("epoch") plt.savefig(file_dir + "/acc.png") plt.close(2)
运行以及结果查看
完成上面的所有代码就可以开始运行了。点击右键,然后选择“run train.py”即可,运行结果如下:
在每个epoch测试完成之后,打印验证集的acc、recall等指标。
DeiT测试结果:
DeiT_dist测试结果:
测试
测试,我们采用一种通用的方式。
测试集存放的目录如下图:
DEiT_demo├─test│ ├─1.jpg│ ├─2.jpg│ ├─3.jpg│ ├ ......└─test.py
import torch.utils.data.distributedimport torchvision.transforms as transformsfrom PIL import Imagefrom torch.autograd import Variableimport osclasses = ('Black-grass', 'Charlock', 'Cleavers', 'Common Chickweed', 'Common wheat', 'Fat Hen', 'Loose Silky-bent', 'Maize', 'Scentless Mayweed', 'Shepherds Purse', 'Small-flowered Cranesbill', 'Sugar beet')transform_test = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.51819474, 0.5250407, 0.4945761], std=[0.24228974, 0.24347611, 0.2530049])])DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")model=torch.load('checkpoints/DEiT/best.pth')model.eval()model.to(DEVICE)path = 'test/'testList = os.listdir(path)for file in testList: img = Image.open(path + file) img = transform_test(img) img.unsqueeze_(0) img = Variable(img).to(DEVICE) out = model(img) # Predict _, pred = torch.max(out.data, 1) print('Image Name:{},predict:{}'.format(file, classes[pred.data.item()]))
测试的主要逻辑:
1、定义类别,这个类别的顺序和训练时的类别顺序对应,一定不要改变顺序!!!!
2、定义transforms,transforms和验证集的transforms一样即可,别做数据增强。
3、 加载model,并将模型放在DEVICE里,
4、循环 读取图片并预测图片的类别,在这里注意,读取图片用PIL库的Image。不要用cv2,transforms不支持。循环里面的主要逻辑:
- 使用Image.open读取图片
- 使用transform_test对图片做归一化和标椎化。
- img.unsqueeze_(0) 增加一个维度,由(3,224,224)变为(1,3,224,224)
- Variable(img).to(DEVICE):将数据放入DEVICE中。
- model(img):执行预测。
- _, pred = torch.max(out.data, 1):获取预测值的最大下角标。
运行结果:
完整的代码
https://download.csdn.net/download/hhhhhhhhhhwwwwwwwwww/87294382