目錄
摘要
建立項目
導入所需要的庫
設定全局參數
圖像預處理
讀取資料
設定模型
設定訓練和驗證
測試
完整代碼：
EfficientNet是谷歌2019年提出的分類模型，自從提出以後這個模型，各大競賽平台常常能看到他的身影，成了霸榜的神器。下圖是EfficientNet—B0模型的網絡結構。
從網絡中可以看出，作者建構了MBConv，結構如下圖：
k對應的卷積核的大小，經過1×1的卷積，然後channel放大4倍，再經過depthwise conv3×3的卷積，然後經過SE子產品後，再經過1×1的卷積，把channel恢複到輸入的大小，最後和上層的輸入融合。
本文簡單介紹一下EfficientNet的網絡結構，主要實戰為主，下面講講如何使用EfficientNet實作貓狗分類，由于本文使用的Loss函數是CrossEntropyLoss，是以隻需更改類别的個數就可以實作多分類。
建立一個圖像分類的項目，data裡面放資料集，dataset檔案夾中自定義資料的讀取方法，這次我不采用預設的讀取方式，太簡單沒啥意思。然後再建立train.py和test.py
在項目的根目錄建立train.py，然後在裡面寫訓練代碼。
首先檢查有沒有安裝EfficientNet的庫，如果沒有安裝則執行pip install efficientnet_pytorch安裝EfficientNet庫，安裝後再導入。
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
from dataset.dataset import DogCat
from torch.autograd import Variable
from efficientnet_pytorch import EfficientNet
#pip install efficientnet_pytorch
設定BatchSize、學習率和epochs，判斷是否有cuda環境，如果沒有設定為cpu。
# 設定全局參數
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不收斂的情況。
# 資料預處理
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([
資料集位址：連結：
https://pan.baidu.com/s/1ZM8vDWEzgscJMnBrZfvQGw提取碼：48c3
将其下載下傳後解壓放到data檔案夾中。資料的目錄如下圖：
然後我們在dataset檔案夾下面建立 __init__.py和dataset.py，在dataset.py檔案夾寫入下面的代碼：
# coding:utf8
import os
from PIL import Image
from torch.utils import data
from torchvision import transforms as T
from sklearn.model_selection import train_test_split
class DogCat(data.Dataset):
def __init__(self, root, transforms=None, train=True, test=False):
"""
主要目标： 擷取所有圖檔的位址，并根據訓練，驗證，測試劃分資料
self.test = test
self.transforms = transforms
imgs = [os.path.join(root, img) for img in os.listdir(root)]
if self.test:
imgs = sorted(imgs, key=lambda x: int(x.split('.')[-2].split('/')[-1]))
else:
imgs = sorted(imgs, key=lambda x: int(x.split('.')[-2]))
self.imgs = imgs
trainval_files, val_files = train_test_split(imgs, test_size=0.3, random_state=42)
if train:
self.imgs = trainval_files
else:
self.imgs = val_files
def __getitem__(self, index):
一次傳回一張圖檔的資料
img_path = self.imgs[index]
label =-1
label = 1 if 'dog' in img_path.split('/')[-1] else 0
data = Image.open(img_path)
data = self.transforms(data)
return data, label
def __len__(self):
return len(self.imgs)
然後我們在train.py調用DogCat讀取資料
dataset_train = DogCat('data/train', transforms=transform, train=True)
dataset_test = DogCat("data/train", transforms=transform_test, train=False)
# 讀取資料
print(dataset_train.imgs)
# 導入資料
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，模型采用efficientnet-B3。更改最後一層的全連接配接，将類别設定為2，然後将模型放到DEVICE。優化器選用Adam。
# 執行個體化模型并且移動到GPU
criterion = nn.CrossEntropyLoss()
model_ft = EfficientNet.from_pretrained('efficientnet-b3')
num_ftrs = model_ft._fc.in_features
model_ft._fc = nn.Linear(num_ftrs, 2)
model_ft.to(DEVICE)
# 選擇簡單暴力的Adam優化器，學習率調低
optimizer = optim.Adam(model_ft.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_ft, DEVICE, train_loader, optimizer, epoch)
val(model_ft, DEVICE, test_loader)
torch.save(model_ft, 'model.pth')
完成上面的代碼後就可以開始訓練，點選run開始訓練，如下圖：
由于我們使用了預訓練模型，是以收斂速度很快。
我介紹兩種常用的測試方式，第一種是通用的，通過自己手動加載資料集然後做預測，具體操作如下：
測試集存放的目錄如下圖：
第一步 定義類别，這個類别的順序和訓練時的類别順序對應，一定不要改變順序！！！！我們在訓練時，cat類别是0，dog類别是1，是以我定義classes為(cat,dog)。
第二步 定義transforms，transforms和驗證集的transforms一樣即可，别做資料增強。
第三步 加載model，并将模型放在DEVICE裡，
第四步 讀取圖檔并預測圖檔的類别，在這裡注意，讀取圖檔用PIL庫的Image。不要用cv2，transforms不支援。
import torchvision.datasets as datasets
classes = ('cat', 'dog')
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()]))
運作結果：
第二種使用我們剛才定義的dataset.py加載測試集。代碼如下：
dataset_test =DogCat('data/test/', transform_test,test=True)
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], classes[pred.data.item()]))
index += 1
train.py
BATCH_SIZE = 32
EPOCHS = 10
test1.py
test2.py