持續創作，加速成長！這是我參與「掘金日新計劃 · 6 月更文挑戰」的第17天，點選檢視活動詳情 @[toc]

MobileVIT實戰

論文地址：http://arxiv.org/abs/2110.02178

官方程式碼：http://github.com/apple/ml-cvnets

本文使用的程式碼來自：http://gitcode.net/mirrors/rwightman/pytorch-image-models，也就是大名鼎鼎的timm。

目前，Transformer已經霸榜計算機視覺各種任務，但是缺點也很明顯就是引數量太大無法用在移動裝置，為了解決這個問題，Apple的科學家們將CNN和VIT的優勢結合起來，提出了一個輕量級的視覺網路模型mobileViT。

根據論文中給出的Top-1成績的對比結果，我們可以得出，xs模型引數量比經典的MobileNetV3小，但是精度卻提高了7.4%,標準的S模型比ResNet-101，還高一些，但是引數量也只有ResNet-101的九分之一。這樣的成績可謂逆天了！

本文從實戰的角度出發，帶領大家感受一下mobileViT，我們還是使用以前的植物分類資料集，模型採用MobileViT-S。

安裝timm

安裝timm，使用pip就行，命令：

python pip install timm

安裝完成之後，才發現沒有MobileViT，我以為是晚上太晚了，眼睛不好使了。後來才發現，pip安裝的最新版本只有0.54，但是官方最新的版本是0.61，所以只能換種方式安裝了。

登入到官方的GitHub，mirrors / rwightman / pytorch-image-models · GitCode，將其下載到本地，然後執行命令：

python python setup.py install

安裝完成後就可以找到mobileViT了。

建議使用timm，因為timm有預訓練，這樣可以加快訓練速度。

資料增強Cutout和Mixup

為了提高成績我在程式碼中加入Cutout和Mixup這兩種增強方式。實現這兩種增強需要安裝torchtoolbox。安裝命令：

pip install torchtoolbox

Cutout實現，在transforms中。

```python from torchtoolbox.transform import Cutout

資料預處理

transform = transforms.Compose([ transforms.Resize((224, 224)), Cutout(), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])

]) ```

需要匯入包：from timm.data.mixup import Mixup，

定義Mixup，和SoftTargetCrossEntropy

```python 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=12)

criterion_train = SoftTargetCrossEntropy() ```

專案結構

MobileVIT_demo ├─data │ ├─Black-grass │ ├─Charlock │ ├─Cleavers │ ├─Common Chickweed │ ├─Common wheat │ ├─Fat Hen │ ├─Loose Silky-bent │ ├─Maize │ ├─Scentless Mayweed │ ├─Shepherds Purse │ ├─Small-flowered Cranesbill │ └─Sugar beet ├─mean_std.py ├─makedata.py ├─train.py └─test.py

mean_std.py：計算mean和std的值。

makedata.py：生成資料集。

計算mean和std

為了使模型更加快速的收斂，我們需要計算出mean和std的值，新建mean_std.py,插入程式碼：

```python from torchvision.datasets import ImageFolder import torch from torchvision import transforms

def get_mean_and_std(train_data): train_loader = torch.utils.data.DataLoader( train_data, batch_size=1, shuffle=False, num_workers=0, pin_memory=True) mean = torch.zeros(3) std = torch.zeros(3) for X, _ in train_loader: for d in range(3): mean[d] += X[:, d, :, :].mean() std[d] += X[:, d, :, :].std() mean.div_(len(train_data)) std.div_(len(train_data)) return list(mean.numpy()), list(std.numpy())

if name == 'main': train_dataset = ImageFolder(root=r'data1', transform=transforms.ToTensor()) print(get_mean_and_std(train_dataset)) ```

資料集結構：

執行結果：

([0.3281186, 0.28937867, 0.20702125], [0.09407319, 0.09732835, 0.106712654])

把這個結果記錄下來，後面要用！

生成資料集

我們整理還的影象分類的資料集結構是這樣的

data ├─Black-grass ├─Charlock ├─Cleavers ├─Common Chickweed ├─Common wheat ├─Fat Hen ├─Loose Silky-bent ├─Maize ├─Scentless Mayweed ├─Shepherds Purse ├─Small-flowered Cranesbill └─Sugar beet

pytorch和keras預設載入方式是ImageNet資料集格式，格式是

├─data │ ├─val │ │ ├─Black-grass │ │ ├─Charlock │ │ ├─Cleavers │ │ ├─Common Chickweed │ │ ├─Common wheat │ │ ├─Fat Hen │ │ ├─Loose Silky-bent │ │ ├─Maize │ │ ├─Scentless Mayweed │ │ ├─Shepherds Purse │ │ ├─Small-flowered Cranesbill │ │ └─Sugar beet │ └─train │ ├─Black-grass │ ├─Charlock │ ├─Cleavers │ ├─Common Chickweed │ ├─Common wheat │ ├─Fat Hen │ ├─Loose Silky-bent │ ├─Maize │ ├─Scentless Mayweed │ ├─Shepherds Purse │ ├─Small-flowered Cranesbill │ └─Sugar beet

新增格式轉化指令碼makedata.py,插入程式碼：

```python import glob import os import shutil

image_list=glob.glob('data1//.png') print(image_list) file_dir='data' if os.path.exists(file_dir): print('true') #os.rmdir(file_dir) shutil.rmtree(file_dir)#刪除再建立 os.makedirs(file_dir) else: os.makedirs(file_dir)

from sklearn.model_selection import train_test_split trainval_files, val_files = train_test_split(image_list, test_size=0.3, random_state=42) train_dir='train' val_dir='val' train_root=os.path.join(file_dir,train_dir) val_root=os.path.join(file_dir,val_dir) for file in trainval_files: file_class=file.replace("\","/").split('/')[-2] file_name=file.replace("\","/").split('/')[-1] file_class=os.path.join(train_root,file_class) if not os.path.isdir(file_class): os.makedirs(file_class) shutil.copy(file, file_class + '/' + file_name)

for file in val_files: file_class=file.replace("\","/").split('/')[-2] file_name=file.replace("\","/").split('/')[-1] file_class=os.path.join(val_root,file_class) if not os.path.isdir(file_class): os.makedirs(file_class) shutil.copy(file, file_class + '/' + file_name) ```

訓練

完成上面的步驟後，就開始train指令碼的編寫，新建train.py.

匯入專案使用的庫

python 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.datasets as datasets import torchvision.transforms as transforms from sklearn.metrics import classification_report from timm.data.mixup import Mixup from timm.loss import SoftTargetCrossEntropy from timm.models.mobilevit import mobilevit_s from apex import amp import warnings warnings.filterwarnings("ignore")

設定全域性引數

設定學習率、BatchSize、epoch等引數，判斷環境中是否存在GPU，如果沒有則使用CPU。建議使用GPU，CPU太慢了。

```python

設定全域性引數

model_lr = 1e-4 BATCH_SIZE = 8 EPOCHS = 300 DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') use_amp=False #是否使用混合精度 classes=12

資料預處理7

```

model_lr：學習率，根據實際情況做調整。

BATCH_SIZE：batchsize，根據顯示卡的大小設定。

EPOCHS：epoch的個數，一般300夠用。

use_amp：是否使用混合精度。

classes：類別個數。

CLIP_GRAD：梯度的最大範數，在梯度裁剪裡設定。

影象預處理與增強

資料處理比較簡單，加入了Cutout、做了Resize和歸一化，定義Mixup函式。

這裡注意下Resize的大小，由於MobileViT的輸入是256×256的大小，所以要Resize為256×256。

```python

資料預處理7

transform = transforms.Compose([ transforms.Resize((256, 256)), Cutout(), transforms.ToTensor(), transforms.Normalize(mean=[0.51819474, 0.5250407, 0.4945761], std=[0.24228974, 0.24347611, 0.2530049])

]) transform_test = transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize(mean=[0.51819474, 0.5250407, 0.4945761], std=[0.24228974, 0.24347611, 0.2530049]) ]) 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檔案中。

```python

讀取資料

dataset_train = datasets.ImageFolder('data/train', transform=transform) dataset_test = datasets.ImageFolder("data/val", transform=transform_test)

匯入資料

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) print(dataset_train.class_to_idx) 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))

```

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}

設定模型

• 設定loss函式，train的loss為：SoftTargetCrossEntropy，val的loss：nn.CrossEntropyLoss()。
• 設定模型為mobilevit_s，預訓練設定為true，num_classes設定為12。
• 優化器設定為adam。
• 學習率調整策略選擇為餘弦退火。
• 檢測可用顯示卡的數量，如果大於1，則要用torch.nn.DataParallel載入模型，開啟多卡訓練。
• 開啟混合精度訓練。
• 如果存在多上顯示卡，則使用DP的方式開啟多卡並行訓練。

```python

例項化模型並且移動到GPU

criterion_train = SoftTargetCrossEntropy()# 訓練用的loss criterion_val = torch.nn.CrossEntropyLoss()# 驗證用的loss model_ft = mobilevit_s(pretrained=True)# 定義模型，並設定預訓練 print(model_ft) num_ftrs = model_ft.head.fc.in_features model_ft.head.fc = nn.Linear(num_ftrs, classes)# 修改類別 model_ft.to(DEVICE) print(model_ft)

選擇簡單暴力的Adam優化器，學習率調低

optimizer = optim.Adam(model_ft.parameters(), lr=model_lr) cosine_schedule = optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer, T_max=20, eta_min=1e-6)# 使用餘弦退火演算法調整學習率 if use_amp: #如果使用混合精度訓練，則初始化amp。 model, optimizer = amp.initialize(model_ft, optimizer, opt_level="O1") # 這裡是“歐一”，不是“零一” if torch.cuda.device_count() > 1: #檢測是否存在多張顯示卡，如果存在則使用DP的方式並行訓練 print("Let's use", torch.cuda.device_count(), "GPUs!") model_ft = torch.nn.DataParallel(model_ft) ```

定義訓練和驗證函式

定義訓練函式和驗證函式，在一個epoch完成後，使用classification_report計算詳細的得分情況。

訓練的主要步驟：

1、判斷迭代的資料是否是奇數，由於mixup_fn只能接受偶數，所以如果不是偶數則要減去一位，讓其變成偶數。但是有可能最後一次迭代只有一條資料，減去後就變成了0，所以還要判斷不能小於2，如果小於2則直接中斷本次迴圈。

2、將資料輸入mixup_fn生成mixup資料，然後輸入model計算loss。

3、如果使用混合精度，則使用amp.scale_loss反向傳播求解梯度，否則，直接反向傳播求梯度。torch.nn.utils.clip_grad_norm_函式執行梯度裁剪，防止梯度爆炸。

等待一個epoch完成後，統計類別的得分情況。

```python

定義訓練過程

def train(model, device, train_loader, optimizer, epoch): model.train() sum_loss = 0 total_num = len(train_loader.dataset) print(total_num, len(train_loader)) for batch_idx, (data, target) in enumerate(train_loader): if len(data) % 2 != 0: if len(data) < 2: continue data = data[0:len(data) - 1] target = target[0:len(target) - 1] print(len(data)) data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True) samples, targets = mixup_fn(data, target) output = model(data) loss = criterion_train(output, targets) optimizer.zero_grad() if use_amp: with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), CLIP_GRAD) else: loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), CLIP_GRAD) optimizer.step() lr = optimizer.state_dict()['param_groups'][0]['lr'] print_loss = loss.data.item() sum_loss += print_loss 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 = sum_loss / len(train_loader) print('epoch:{},loss:{}'.format(epoch, ave_loss))

ACC = 0

驗證過程

def val(model, device, test_loader): global ACC model.eval() test_loss = 0 correct = 0 total_num = len(test_loader.dataset) print(total_num, len(test_loader)) val_list = [] pred_list = [] with torch.no_grad(): for data, target in test_loader: for t in target: val_list.append(t.data.item()) data, target = data.to(device), target.to(device) output = model(data) loss = criterion_val(output, target) , pred = torch.max(output.data, 1) for p in pred: pred_list.append(p.data.item()) correct += torch.sum(pred == target) print_loss = loss.data.item() test_loss += print_loss correct = correct.data.item() acc = correct / total_num avgloss = test_loss / len(test_loader) print('\nVal set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( avgloss, correct, len(test_loader.dataset), 100 * acc)) if acc > ACC: if isinstance(model, torch.nn.DataParallel): torch.save(model.module, 'model' + str(epoch) + '' + str(round(acc, 3)) + '.pth') else: torch.save(model, 'model' + str(epoch) + '_' + str(round(acc, 3)) + '.pth') ACC = acc return val_list, pred_list

訓練

is_set_lr = False for epoch in range(1, EPOCHS + 1): train(model_ft, DEVICE, train_loader, optimizer, epoch) if epoch < 600: cosine_schedule.step() else: if is_set_lr: continue for param_group in optimizer.param_groups: param_group["lr"] = 1e-6 is_set_lr = True val_list, pred_list = val(model_ft, DEVICE, test_loader) print(classification_report(val_list, pred_list, target_names=dataset_train.class_to_idx))

```

執行結果：

測試

測試，我們採用一種通用的方式。

測試集存放的目錄如下圖：

第一步 定義類別，這個類別的順序和訓練時的類別順序對應，一定不要改變順序！！！！

第二步 定義transforms，transforms和驗證集的transforms一樣即可，別做資料增強。

第三步 載入model，並將模型放在DEVICE裡，

第四步 讀取圖片並預測圖片的類別，在這裡注意，讀取圖片用PIL庫的Image。不要用cv2，transforms不支援。

```python import torch.utils.data.distributed import torchvision.transforms as transforms from PIL import Image from torch.autograd import Variable import os

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((256, 256)), 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("model_52_0.954.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()])) ```

執行結果：

完整程式碼：

http://download.csdn.net/download/hhhhhhhhhhwwwwwwwwww/85232437