【影象分類】實戰——使用EfficientNetV2實現影象分類（Pytorch）

語言: CN / TW / HK
時間 2022-06-28 05:15:30

目錄

摘要

新建專案

匯入所需要的庫

設定全域性引數

影象預處理

讀取資料

設定模型

設定訓練和驗證

驗證

完整程式碼：

摘要

這幾天學習了EfficientNetV2，對論文做了翻譯，並復現了論文的程式碼。

論文翻譯：【影象分類】 EfficientNetV2：更快、更小、更強——論文翻譯_AI浩-CSDN部落格

程式碼復現：【影象分類】用通俗易懂程式碼的復現EfficientNetV2，入門的絕佳選擇（pytorch）_AI浩-CSDN部落格

對EfficientNetV2想要了解的可以檢視上面的文章，這篇文章著重介紹如何使用EfficientNetV2實現影象分類。Loss函式採用CrossEntropyLoss,可以通過更改最後一層的全連線方便實現二分類和多分類。資料集採用經典的貓狗大戰資料集，做二分類的實現。

資料集地址：連結：https\://pan.baidu.com/s/1kqhVPOqV5vklYYIFVAzAAA \ 提取碼：3ch6 

新建專案

新建一個影象分類的專案，在專案的跟目錄新建資料夾model，用於存放EfficientNetV2的模型程式碼，新建EfficientNetV2.py，將【影象分類】用通俗易懂程式碼的復現EfficientNetV2，入門的絕佳選擇（pytorch）_AI浩-CSDN部落格 復現的程式碼複製到裡面，然後在model資料夾新建__init__.py空檔案，model的目錄結構如下：

在專案的根目錄新建train.py，然後在裡面寫訓練程式碼。

匯入所需要的庫

import torch.optim as optim import torch import torch.nn as nn import torch.nn.parallel import torch.optim import torch.utils.data import torch.utils.data.distributed import torchvision.transforms as transforms import torchvision.datasets as datasets from torch.autograd import Variable from model.EfficientNetv2 import efficientnetv2_s

設定全域性引數

\ 設定BatchSize、學習率和epochs，判斷是否有cuda環境，如果沒有設定為cpu。

```python

設定全域性引數

modellr = 1e-4 BATCH_SIZE = 64 EPOCHS = 20 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') ```

影象預處理

在做影象與處理時，train資料集的transform和驗證集的transform分開做，train的影象處理出了resize和歸一化之外，還可以設定影象的增強，比如旋轉、隨機擦除等一系列的操作，驗證集則不需要做影象增強，另外不要盲目的做增強，不合理的增強手段很可能會帶來負作用，甚至出現Loss不收斂的情況。

```python

資料預處理

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

讀取資料

使用Pytorch的預設方式讀取資料。資料的目錄如下圖：

​

訓練集，取了貓狗大戰資料集中，貓狗影象各一萬張，剩餘的放到驗證集中。

```python

讀取資料

dataset_train = datasets.ImageFolder('data/train', transform) print(dataset_train.imgs)

對應資料夾的label

print(dataset_train.class_to_idx) dataset_test = datasets.ImageFolder('data/val', transform_test)

對應資料夾的label

print(dataset_test.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) ```

設定模型

\ 使用CrossEntropyLoss作為loss，模型採用efficientnetv2_s，由於沒有Pytorch的預訓練模型，我們只能從頭開始訓練。更改最後一層的全連線，將類別設定為2，然後將模型放到DEVICE。優化器選用Adam。

```python

例項化模型並且移動到GPU

criterion = nn.CrossEntropyLoss() model = efficientnetv2_s() num_ftrs = model.classifier.in_features model.classifier = nn.Linear(num_ftrs, 2) model.to(DEVICE)

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

optimizer = optim.Adam(model.parameters(), lr=modellr)

def adjust_learning_rate(optimizer, epoch): """Sets the learning rate to the initial LR decayed by 10 every 30 epochs""" modellrnew = modellr * (0.1 ** (epoch // 50)) print("lr:", modellrnew) for param_group in optimizer.param_groups: param_group['lr'] = modellrnew ```

設定訓練和驗證

```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): data, target = Variable(data).to(device), Variable(target).to(device) output = model(data) loss = criterion(output, target) optimizer.zero_grad() loss.backward() optimizer.step() print_loss = loss.data.item() sum_loss += print_loss if (batch_idx + 1) % 50 == 0: print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, (batch_idx + 1) * len(data), len(train_loader.dataset), 100. * (batch_idx + 1) / len(train_loader), loss.item())) ave_loss = sum_loss / len(train_loader) print('epoch:{},loss:{}'.format(epoch, ave_loss))

驗證過程

def val(model, device, test_loader): model.eval() test_loss = 0 correct = 0 total_num = len(test_loader.dataset) print(total_num, len(test_loader)) with torch.no_grad(): for data, target in test_loader: data, target = Variable(data).to(device), Variable(target).to(device) output = model(data) loss = criterion(output, target) _, pred = torch.max(output.data, 1) 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))

訓練

for epoch in range(1, EPOCHS + 1): adjust_learning_rate(optimizer, epoch) train(model, DEVICE, train_loader, optimizer, epoch) val(model, DEVICE, test_loader) torch.save(model, 'model.pth') ```

​完成上面的程式碼後就可以開始訓練，點選run開始訓練，如下圖：

驗證

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

第一步 定義類別，這個類別的順序和訓練時的類別順序對應，一定不要改變順序！！！！我們在訓練時，cat類別是0，dog類別是1，所以我定義classes為(cat,dog)。

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

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

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

```python import torch.utils.data.distributed import torchvision.transforms as transforms

from torch.autograd import Variable import os from PIL import Image

classes = ('cat', 'dog')

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

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model = torch.load("model.pth") model.eval() model.to(DEVICE) path='data/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()]))

```

\  \ 執行結果：

其實在讀取資料，也可以巧妙的用datasets.ImageFolder，下面我們就用datasets.ImageFolder實現對圖片的預測。改一下test資料集的路徑，在test資料夾外面再加一層檔案件，取名為dataset，如下圖所示：

然後修改讀取圖片的方式。程式碼如下：

```python import torch.utils.data.distributed import torchvision.transforms as transforms import torchvision.datasets as datasets from torch.autograd import Variable   classes = ('cat', 'dog') transform_test = transforms.Compose([     transforms.Resize((224, 224)),     transforms.ToTensor(),     transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) ])   DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model = torch.load("model.pth") model.eval() model.to(DEVICE)   dataset_test = datasets.ImageFolder('data/datatest', transform_test) print(len(dataset_test))

對應資料夾的label

for index in range(len(dataset_test)):     item = dataset_test[index]     img, label = item     img.unsqueeze_(0)     data = Variable(img).to(DEVICE)     output = model(data)     _, pred = torch.max(output.data, 1)     print('Image Name:{},predict:{}'.format(dataset_test.imgs[index][0], classes[pred.data.item()]))     index += 1   ```

完整程式碼：

train.py

```python import torch.optim as optim import torch import torch.nn as nn import torch.nn.parallel import torch.optim import torch.utils.data import torch.utils.data.distributed import torchvision.transforms as transforms import torchvision.datasets as datasets from torch.autograd import Variable from model.EfficientNetv2 import efficientnetv2_s

設定全域性引數

modellr = 1e-4 BATCH_SIZE = 32 EPOCHS = 50 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

資料預處理

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

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

讀取資料

dataset_train = datasets.ImageFolder('data/train', transform) print(dataset_train.imgs)

對應資料夾的label

print(dataset_train.class_to_idx) dataset_test = datasets.ImageFolder('data/val', transform_test)

對應資料夾的label

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

例項化模型並且移動到GPU

criterion = nn.CrossEntropyLoss() model = efficientnetv2_s() num_ftrs = model.classifier.in_features model.classifier = nn.Linear(num_ftrs, 2) model.to(DEVICE)

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

optimizer = optim.Adam(model.parameters(), lr=modellr)

def adjust_learning_rate(optimizer, epoch): """Sets the learning rate to the initial LR decayed by 10 every 30 epochs""" modellrnew = modellr * (0.1 ** (epoch // 50)) print("lr:", modellrnew) for param_group in optimizer.param_groups: param_group['lr'] = modellrnew

定義訓練過程

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): data, target = Variable(data).to(device), Variable(target).to(device) output = model(data) loss = criterion(output, target) optimizer.zero_grad() loss.backward() optimizer.step() print_loss = loss.data.item() sum_loss += print_loss if (batch_idx + 1) % 50 == 0: print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, (batch_idx + 1) * len(data), len(train_loader.dataset), 100. * (batch_idx + 1) / len(train_loader), loss.item())) ave_loss = sum_loss / len(train_loader) print('epoch:{},loss:{}'.format(epoch, ave_loss))

驗證過程

def val(model, device, test_loader): model.eval() test_loss = 0 correct = 0 total_num = len(test_loader.dataset) print(total_num, len(test_loader)) with torch.no_grad(): for data, target in test_loader: data, target = Variable(data).to(device), Variable(target).to(device) output = model(data) loss = criterion(output, target) _, pred = torch.max(output.data, 1) 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))

訓練

for epoch in range(1, EPOCHS + 1): adjust_learning_rate(optimizer, epoch) train(model, DEVICE, train_loader, optimizer, epoch) val(model, DEVICE, test_loader) torch.save(model, 'model.pth') ```

test1.py

```python

import torch.utils.data.distributed import torchvision.transforms as transforms

from torch.autograd import Variable import os from PIL import Image

classes = ('cat', 'dog')

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

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model = torch.load("model.pth") model.eval() model.to(DEVICE) path='data/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()]))

```

test2.py

``` import torch.utils.data.distributed import torchvision.transforms as transforms import torchvision.datasets as datasets from torch.autograd import Variable   classes = ('cat', 'dog') transform_test = transforms.Compose([     transforms.Resize((224, 224)),     transforms.ToTensor(),     transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) ])   DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model = torch.load("model.pth") model.eval() model.to(DEVICE)   dataset_test = datasets.ImageFolder('data/datatest', transform_test) print(len(dataset_test))

對應資料夾的label

for index in range(len(dataset_test)):     item = dataset_test[index]     img, label = item     img.unsqueeze_(0)     data = Variable(img).to(DEVICE)     output = model(data)     _, pred = torch.max(output.data, 1)     print('Image Name:{},predict:{}'.format(dataset_test.imgs[index][0], classes[pred.data.item()]))     index += 1 ```

「」