CloFormer實戰(zhàn)：使用CloFormer實現(xiàn)圖像分類任務(wù)（二）
在上一篇文章中完成了前期的準備工作，見鏈接：CloFormer實戰(zhàn)：使用CloFormer實現(xiàn)圖像分類任務(wù)（一）這篇主要是講解如何訓練和測試
訓練部分
完成上面的步驟后，就開始train腳本的編寫，新建train.py
導入項目使用的庫
在train.py導入
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?models.model_factory?import?cloformer_xxs
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" 選擇顯卡，index從0開始，比如一臺機器上有8塊顯卡，我們打算使用前兩塊顯卡訓練，設(shè)置為“0,1”,同理如果打算使用第三塊和第六塊顯卡訓練，則設(shè)置為“2,5”。
設(shè)置隨機因子
def?seed_everything(seed=42):
????os.environ['PYHTONHASHSEED']?=?str(seed)
????torch.manual_seed(seed)
????torch.cuda.manual_seed(seed)
????torch.backends.cudnn.deterministic?=?True
設(shè)置了固定的隨機因子，再次訓練的時候就可以保證圖片的加載順序不會發(fā)生變化。
設(shè)置全局參數(shù)
if?__name__?==?'__main__':
????#創(chuàng)建保存模型的文件夾
????file_dir?=?'checkpoints/cloformer/'
????if?os.path.exists(file_dir):
????????print('true')
????????os.makedirs(file_dir,exist_ok=True)
????else:
????????os.makedirs(file_dir)
????#?設(shè)置全局參數(shù)
????model_lr?=?1e-4
????BATCH_SIZE?=?16
????EPOCHS?=?300
????DEVICE?=?torch.device('cuda:0'?if?torch.cuda.is_available()?else?'cpu')
????use_amp?=?True??#?是否使用混合精度
????use_dp?=?True?#是否開啟dp方式的多卡訓練
????classes?=?12
????resume?="checkpoints/cloformer/model_2_44.606.pth"
????CLIP_GRAD?=?5.0
????Best_ACC?=?0?#記錄最高得分
????use_ema=True
????model_ema_decay=0.9998
????start_epoch=1
????seed=1
????seed_everything(seed)
設(shè)置存放權(quán)重文件的文件夾，如果文件夾存在刪除再建立。
接下來，設(shè)置全局參數(shù)，比如：學習率、BatchSize、epoch等參數(shù)，判斷環(huán)境中是否存在GPU，如果沒有則使用CPU。 注：建議使用GPU，CPU太慢了。
參數(shù)的詳細解釋：
★model_lr：學習率，根據(jù)實際情況做調(diào)整。
”
BATCH_SIZE：batchsize，根據(jù)顯卡的大小設(shè)置。
EPOCHS：epoch的個數(shù)，一般300夠用。
use_amp：是否使用混合精度。
use_dp ：是否開啟dp方式的多卡訓練？
classes：類別個數(shù)。
resume：再次訓練的模型路徑，如果不為None,則表示加載resume指向的模型繼續(xù)訓練。
CLIP_GRAD：梯度的最大范數(shù)，在梯度裁剪里設(shè)置。
Best_ACC：記錄最高ACC得分。
use_ema：是否使用ema
start_epoch：開始的epoch，默認是1，如果重新訓練時，需要給start_epoch重新賦值。
SEED:隨機因子，數(shù)值可以隨意設(shè)定，但是設(shè)置后，不要隨意更改，更改后，圖片加載的順序會改變，影響測試結(jié)果。
?file_dir?=?'checkpoints/cloformer'
這是存放seaformer模型的路徑。
圖像預處理與增強
???#?數(shù)據(jù)預處理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.3281186,?0.28937867,?0.20702125],?std=?[0.09407319,?0.09732835,?0.106712654])
????])
????transform_test?=?transforms.Compose([
????????transforms.Resize((224,?224)),
????????transforms.ToTensor(),
????????transforms.Normalize(mean=[0.3281186,?0.28937867,?0.20702125],?std=?[0.09407319,?0.09732835,?0.106712654])
????])
????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)
數(shù)據(jù)處理和增強比較簡單，加入了隨機10度的旋轉(zhuǎn)、高斯模糊、色彩飽和度明亮度的變化、Mixup等比較常用的增強手段，做了Resize和歸一化。
?transforms.Normalize(mean=[0.3281186,?0.28937867,?0.20702125],?std=?[0.09407319,?0.09732835,?0.106712654])
這里設(shè)置為計算mean和std。 這里注意下Resize的大小，由于選用的CloFormer模型輸入是224×224的大小，所以要Resize為224×224。
讀取數(shù)據(jù)
???#?讀取數(shù)據(jù)
????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))
????#?導入數(shù)據(jù)
????train_loader?=?torch.utils.data.DataLoader(dataset_train,?batch_size=BATCH_SIZE,?pin_memory=True,shuffle=True,drop_last=True)
????test_loader?=?torch.utils.data.DataLoader(dataset_test,?batch_size=BATCH_SIZE,?pin_memory=True,shuffle=False)
使用pytorch默認讀取數(shù)據(jù)的方式，然后將dataset_train.class_to_idx打印出來，預測的時候要用到。
對于train_loader ，drop_last設(shè)置為True，因為使用了Mixup數(shù)據(jù)增強，必須保證每個batch里面的圖片個數(shù)為偶數(shù)（不能為零），如果最后一個batch里面的圖片為奇數(shù)，則會報錯，所以舍棄最后batch的迭代。pin_memory設(shè)置為True，可以加快運行速度。
將dataset_train.class_to_idx保存到txt文件或者json文件中。
class_to_idx的結(jié)果：
{'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}
設(shè)置Loss
??#?實例化模型并且移動到GPU
????criterion_train?=?SoftTargetCrossEntropy()
????criterion_val?=?torch.nn.CrossEntropyLoss()
設(shè)置loss函數(shù)，訓練的loss為：SoftTargetCrossEntropy，驗證的loss：nn.CrossEntropyLoss()。
設(shè)置模型
??#設(shè)置模型
????model_ft?=?cloformer_xxs()
????num_fr=model_ft.head.in_features
????model_ft.head=nn.Linear(num_fr,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)
????print(model_ft)
設(shè)置模型為cloformer_xxs。將model的head的輸出修改為classes的長度。
如果resume為True，則加載模型接著resume指向的模型接著訓練，使用模型里的Best_ACC初始化Best_ACC，使用epoch參數(shù)初始化start_epoch。
設(shè)置優(yōu)化器和學習率調(diào)整算法
???#?選擇簡單暴力的Adam優(yōu)化器，學習率調(diào)低
???optimizer?=?optim.AdamW(model_ft.parameters(),lr=model_lr)
???cosine_schedule?=?optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer,?T_max=20,?eta_min=1e-6)
優(yōu)化器設(shè)置為adamW。
學習率調(diào)整策略選擇為余弦退火。
設(shè)置混合精度，DP多卡，EMA
????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:
????????model_ema?=?ModelEma(
????????????model_ft,
????????????decay=model_ema_decay,
????????????device=DEVICE,
????????????resume=resume)
????else:
????????model_ema=None
use_amp為True，則開啟混合精度訓練,聲明pytorch自帶的混合精度 torch.cuda.amp.GradScaler()。
檢測可用顯卡的數(shù)量，如果大于1，并且開啟多卡訓練的情況下，則要用torch.nn.DataParallel加載模型，開啟多卡訓練。
如果使用ema，則注冊ema 注：torch.nn.DataParallel方式，默認不能開啟混合精度訓練的，如果想要開啟混合精度訓練，則需要在模型的forward前面加上@autocast()函數(shù)。
如果不開啟混合精度則要將@autocast()去掉，否則loss一直試nan。
定義訓練和驗證函數(shù)
訓練函數(shù)
#?定義訓練過程
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
訓練的主要步驟：
★1、使用AverageMeter保存自定義變量，包括loss，ACC1，ACC5。
2、進入循環(huán)，將data和target放入device上，non_blocking設(shè)置為True。如果pin_memory=True的話，將數(shù)據(jù)放入GPU的時候，也應該把non_blocking打開，這樣就只把數(shù)據(jù)放入GPU而不取出，訪問時間會大大減少。 如果pin_memory=False時，則將non_blocking設(shè)置為False。
3、將數(shù)據(jù)輸入mixup_fn生成mixup數(shù)據(jù)。
4、將第三部生成的mixup數(shù)據(jù)輸入model，輸出預測結(jié)果，然后再計算loss。
5、 optimizer.zero_grad() 梯度清零，把loss關(guān)于weight的導數(shù)變成0。
6、如果使用混合精度，則★”
with torch.cuda.amp.autocast(),開啟混合精度。
計算loss。torch.nan_to_num將輸入中的NaN、正無窮大和負無窮大替換為NaN、posinf和neginf。默認情況下，nan會被替換為零，正無窮大會被替換為輸入的dtype所能表示的最大有限值，負無窮大會被替換為輸入的dtype所能表示的最小有限值。
scaler.scale(loss).backward()，梯度放大。
torch.nn.utils.clip_grad_norm_，梯度裁剪，放置梯度爆炸。
scaler.step(optimizer) ，首先把梯度值unscale回來，如果梯度值不是inf或NaN,則調(diào)用optimizer.step()來更新權(quán)重，否則，忽略step調(diào)用，從而保證權(quán)重不更新。
更新下一次迭代的scaler。
否則，直接反向傳播求梯度。torch.nn.utils.clip_grad_norm_函數(shù)執(zhí)行梯度裁剪，防止梯度爆炸。
”
7、如果use_ema為True，則執(zhí)行model_ema的updata函數(shù)，更新模型。
8、 torch.cuda.synchronize()，等待上面所有的操作執(zhí)行完成。
9、接下來，更新loss，ACC1，ACC5的值。
等待一個epoch訓練完成后，計算平均loss和平均acc
驗證函數(shù)
#?驗證過程
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
驗證集和訓練集大致相似，主要步驟：
★1、在val的函數(shù)上面添加@torch.no_grad()，作用：所有計算得出的tensor的requires_grad都自動設(shè)置為False。即使一個tensor（命名為x）的requires_grad = True，在with torch.no_grad計算，由x得到的新tensor（命名為w-標量）requires_grad也為False，且grad_fn也為None,即不會對w求導。
2、定義參數(shù): loss_meter: 測試的loss acc1_meter:top1的ACC。 acc5_meter：top5的ACC。 total_num：總的驗證集的數(shù)量。 val_list：驗證集的label。 pred_list：預測的label。
3、進入循環(huán)，迭代test_loader：★將label保存到val_list。
”
將data和target放入device上，non_blocking設(shè)置為True。
將data輸入到model中，求出預測值，然后輸入到loss函數(shù)中，求出loss。
調(diào)用torch.max函數(shù)，將預測值轉(zhuǎn)為對應的label。
將輸出的預測值的label存入pred_list。
調(diào)用accuracy函數(shù)計算ACC1和ACC5
更新loss_meter、acc1_meter、acc5_meter的參數(shù)。4、本次epoch循環(huán)完成后，求得本次epoch的acc、loss。 5、接下來是保存模型的邏輯 如果ACC比Best_ACC高，則保存best模型 判斷模型是否為DP方式訓練的模型。
★如果是DP方式訓練的模型，模型參數(shù)放在model.module，則需要保存model.module。 否則直接保存model。 注：保存best模型，我們采用保存整個模型的方式，這樣保存的模型包含網(wǎng)絡(luò)結(jié)構(gòu)，在預測的時候，就不用再重新定義網(wǎng)絡(luò)了。
”6、接下來保存每個epoch的模型。 判斷模型是否為DP方式訓練的模型。
★如果是DP方式訓練的模型，模型參數(shù)放在model.module，則需要保存model.module.state_dict()。
”
新建個字典，放置Best_ACC、epoch和 model.module.state_dict()等參數(shù)。然后將這個字典保存。判斷是否是使用EMA，如果使用，則還需要保存一份ema的權(quán)重。 否則，新建個字典，放置Best_ACC、epoch和 model.state_dict()等參數(shù)。然后將這個字典保存。判斷是否是使用EMA，如果使用，則還需要保存一份ema的權(quán)重。注意：對于每個epoch的模型只保存了state_dict參數(shù)，沒有保存整個模型文件。
”
調(diào)用訓練和驗證方法
????#?訓練與驗證
????is_set_lr?=?False
????log_dir?=?{}
????train_loss_list,?val_loss_list,?train_acc_list,?val_acc_list,?epoch_list?=?[],?[],?[],?[],?[]
????if?resume?and?os.path.isfile(file_dir+"result.json"):
????????with?open(file_dir+'result.json',?'r',?encoding='utf-8')?as?file:
????????????logs?=?json.load(file)
????????????train_acc_list?=?logs['train_acc']
????????????train_loss_list?=?logs['train_loss']
????????????val_acc_list?=?logs['val_acc']
????????????val_loss_list?=?logs['val_loss']
????????????epoch_list?=?logs['epoch_list']
????for?epoch?in?range(start_epoch,?EPOCHS?+?1):
????????epoch_list.append(epoch)
????????log_dir['epoch_list']?=?epoch_list
????????train_loss,?train_acc?=?train(model_ft,?DEVICE,?train_loader,?optimizer,?epoch,model_ema)
????????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
????????if?use_ema:
????????????val_list,?pred_list,?val_loss,?val_acc?=?val(model_ema.ema,?DEVICE,?test_loader)
????????else:
????????????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)
調(diào)用訓練函數(shù)和驗證函數(shù)的主要步驟：
★1、定義參數(shù)：
is_set_lr，是否已經(jīng)設(shè)置了學習率，當epoch大于一定的次數(shù)后，會將學習率設(shè)置到一定的值，并將其置為True。
log_dir：記錄log用的，將有用的信息保存到字典中，然后轉(zhuǎn)為json保存起來。
train_loss_list：保存每個epoch的訓練loss。
val_loss_list：保存每個epoch的驗證loss。
train_acc_list：保存每個epoch的訓練acc。
val_acc_list：保存么每個epoch的驗證acc。
epoch_list：存放每個epoch的值。
如果是接著上次的斷點繼續(xù)訓練則讀取log文件，然后把log取出來，賦值到對應的list上。 循環(huán)epoch
★”1、調(diào)用train函數(shù)，得到 train_loss, train_acc，并將分別放入train_loss_list，train_acc_list，然后存入到logdir字典中。
2、調(diào)用驗證函數(shù)，判斷是否使用EMA？ 如果使用EMA，則傳入model_ema.ema，否則，傳入model_ft。得到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，將學習率設(shè)置為固定的1e-6。
6、繪制loss曲線和acc曲線。
”
運行以及結(jié)果查看
完成上面的所有代碼就可以開始運行了。點擊右鍵，然后選擇“run train.py”即可，運行結(jié)果如下：
在每個epoch測試完成之后，打印驗證集的acc、recall等指標。
CloFormer測試結(jié)果：
測試
測試，我們采用一種通用的方式。
測試集存放的目錄如下圖：
CloFormer_demo
├─test
│??├─1.jpg
│??├─2.jpg
│??├─3.jpg
│??├?......
└─test.pyimport?torchvision.transforms?as?transforms
from?PIL?import?Image
from?torch.autograd?import?Variable
import?os
import?torch
import?torch.nn?as?nn
classes?=?('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/cloformer/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、定義類別，這個類別的順序和訓練時的類別順序?qū)?，一定不要改變順序！?。。?br>2、定義transforms，transforms和驗證集的transforms一樣即可，別做數(shù)據(jù)增強。
3、 加載model，switch_to_deploy函數(shù)，切換成推理模式，進一步提高運行速度，然后將模型放在DEVICE里，
4、循環(huán) 讀取圖片并預測圖片的類別，在這里注意，讀取圖片用PIL庫的Image。不要用CV2，transforms不支持。循環(huán)里面的主要邏輯：★””
使用Image.open讀取圖片
使用transform_test對圖片做歸一化和標椎化。
img.unsqueeze_(0) 增加一個維度，由（3，224，224）變?yōu)椋?，3，224，224）
Variable(img).to(DEVICE)：將數(shù)據(jù)放入DEVICE中。
model(img)：執(zhí)行預測。
_, pred = torch.max(out.data, 1):獲取預測值的最大下角標。
運行結(jié)果：
熱力圖可視化展示
新建腳本cam_image.py,插入如下代碼：
import?argparse
import?os
import?CV2
import?numpy?as?np
import?torch
from?pytorch_grad_cam?import?GradCAM,?\
????ScoreCAM,?\
????GradCAMPlusPlus,?\
????AblationCAM,?\
????XGradCAM,?\
????EigenCAM,?\
????EigenGradCAM,?\
????LayerCAM,?\
????FullGrad
from?pytorch_grad_cam?import?GuidedBackpropReLUModel
from?pytorch_grad_cam.utils.image?import?show_cam_on_image,?\
????deprocess_image,?\
????preprocess_image
from?pytorch_grad_cam.utils.model_targets?import?ClassifierOutputTarget
import?timm
from?torch.autograd?import?Variable
def?reshape_transform_resmlp(tensor,?height=14,?width=14):
????result?=?tensor.reshape(tensor.size(0),
????????????????????????????height,?width,?tensor.size(2))
????result?=?result.transpose(2,?3).transpose(1,?2)
????return?result
def?reshape_transform_swin(tensor,?height=7,?width=7):
????result?=?tensor.reshape(tensor.size(0),
????????????????????????????height,?width,?tensor.size(2))
????#?Bring?the?channels?to?the?first?dimension,
????#?like?in?CNNs.
????result?=?result.transpose(2,?3).transpose(1,?2)
????return?result
def?reshape_transform_vit(tensor,?height=14,?width=14):
????result?=?tensor[:,?1:,?:].reshape(tensor.size(0),
??????????????????????????????????????height,?width,?tensor.size(2))
????#?Bring?the?channels?to?the?first?dimension,
????#?like?in?CNNs.
????result?=?result.transpose(2,?3).transpose(1,?2)
????return?result
def?get_args():
????parser?=?argparse.ArgumentParser()
????parser.add_argument('--use-cuda',?action='store_true',?default=False,
????????????????????????help='Use?NVIDIA?GPU?acceleration')
????parser.add_argument(
????????'--image-path',
????????type=str,
????????default="./test/0a64e3e6c.png",
????????help='Input?image?path')
????parser.add_argument(
????????'--output-image-path',
????????type=str,
????????default=None,
????????help='Output?image?path')
????parser.add_argument(
????????'--model',
????????type=str,
????????default='cloformer',
????????help='model?name')
????parser.add_argument('--aug_smooth',?action='store_true',
????????????????????????help='Apply?test?time?augmentation?to?smooth?the?CAM')
????parser.add_argument(
????????'--eigen_smooth',
????????action='store_true',
????????help='Reduce?noise?by?taking?the?first?principle?componenet'
?????????????'of?cam_weights*activations')
????parser.add_argument('--method',?type=str,?default='gradcam++',
????????????????????????choices=['gradcam',?'gradcam++',
?????????????????????????????????'scorecam',?'xgradcam',
?????????????????????????????????'ablationcam',?'eigencam',
?????????????????????????????????'eigengradcam',?'layercam',?'fullgrad'],
????????????????????????help='Can?be?gradcam/gradcam++/scorecam/xgradcam'
?????????????????????????????'/ablationcam/eigencam/eigengradcam/layercam')
????args?=?parser.parse_args()
????args.use_cuda?=?args.use_cuda?and?torch.cuda.is_available()
????if?args.use_cuda:
????????print('Using?GPU?for?acceleration')
????else:
????????print('Using?CPU?for?computation')
????return?args
if?__name__?==?'__main__':
????args?=?get_args()
????methods?=?\
????????{"gradcam":?GradCAM,
?????????"scorecam":?ScoreCAM,
?????????"gradcam++":?GradCAMPlusPlus,
?????????"ablationcam":?AblationCAM,
?????????"xgradcam":?XGradCAM,
?????????"eigencam":?EigenCAM,
?????????"eigengradcam":?EigenGradCAM,
?????????"layercam":?LayerCAM,
?????????"fullgrad":?FullGrad}
????DEVICE?=?torch.device("cuda:0"?if?torch.cuda.is_available()?else?"cpu")
????model?=?torch.load('checkpoints/cloformer/best.pth',?map_location='cpu')
????print(model)
????reshape_transform?=?None
????if?'cloformer'?in?args.model:
????????target_layers?=?[model.layers[-1]]??#?[model.network[-1][-2]]
????print(target_layers)
????model.eval()
????model.to(DEVICE)
????img_path?=?args.image_path
????if?args.image_path:
????????img_path?=?args.image_path
????else:
????????import?requests
????????image_url?=?'http://146.48.86.29/edge-mac/imgs/n02123045/ILSVRC2012_val_00023779.JPEG'
????????img_path?=?image_url.split('/')[-1]
????????if?os.path.exists(img_path):
????????????img_data?=?requests.get(image_url).content
????????????with?open(img_path,?'wb')?as?handler:
????????????????handler.write(img_data)
????if?args.output_image_path:
????????save_name?=?args.output_image_path
????else:
????????img_type?=?img_path.split('.')[-1]
????????it_len?=?len(img_type)
????????save_name?=?img_path.split('/')[-1][:-(it_len?+?1)]
????????save_name?=?save_name?+?'_'?+?args.model?+?'.'?+?img_type
????img?=?CV2.imread(img_path,?1)
????img?=?CV2.resize(img,?(224,?224),?interpolation=CV2.INTER_AREA)
????if?args.model?==?'resize':
????????CV2.imwrite(save_name,?img)
????else:
????????rgb_img?=?img[:,?:,?::-1]
????????rgb_img?=?np.float32(rgb_img)?/?255
????????input_tensor?=?Variable(preprocess_image(rgb_img,
?????????????????????????????????????????????????mean=[0.485,?0.456,?0.406],
?????????????????????????????????????????????????std=[0.229,?0.224,?0.225]),?requires_grad=True).to(DEVICE)
????????targets?=?None
????????cam_algorithm?=?methods[args.method]
????????with?cam_algorithm(model=model,
???????????????????????????target_layers=target_layers,
???????????????????????????use_cuda=args.use_cuda,
???????????????????????????reshape_transform=reshape_transform,
???????????????????????????)?as?cam:
????????????cam.batch_size?=?1
????????????grayscale_cam?=?cam(input_tensor=input_tensor,
????????????????????????????????targets=targets,
????????????????????????????????aug_smooth=args.aug_smooth,
????????????????????????????????eigen_smooth=args.eigen_smooth)
????????????grayscale_cam?=?grayscale_cam[0,?:]
????????????cam_image?=?show_cam_on_image(rgb_img,?grayscale_cam,?use_rgb=True)
????????????cam_image?=?CV2.cvtColor(cam_image,?CV2.COLOR_RGB2BGR)
????????CV2.imwrite(save_name,?cam_image)
對get_args函數(shù)的參數(shù)進行設(shè)置：
use-cuda:是否使用cuda，如果在沒有GPU的電腦上調(diào)試時，將其設(shè)置為False。
image-path：待測圖片的路徑，這個是必填項。
model：必填項，默認值：cloformer。
效果如下圖所示：
完整的代碼
完整的代碼： https://download.csdn.net/download/hhhhhhhhhhwwwwwwwwww/87902659
