前段时间参加kaggle的图像隐写的比赛,第一次参加,有幸拿了银牌,使用的就是Efficient这个网络,这个网络在分类任务上确实好用,在此做个简单的记录,如何使用Efficient完成简单的分类任务。
如果有错误的地方,欢迎大家指正
如下图所示,data文件是项目的数据集,Type1~Type4文件夹代表不同类别的图片
首先看第一个代码块“dataset.py”
"""
dataset.py
"""
from glob import glob
from sklearn.model_selection import GroupKFold
import torch
import os
import random
import cv2
import pandas as pd
import numpy as np
import albumentations as A
from albumentations.pytorch.transforms import ToTensorV2
from torch.utils.data import Dataset, DataLoader
import warnings
warnings.filterwarnings("ignore") # 避免烦人的各类无用报警
dataset = []
DATA_ROOT_PATH = "data" # 数据集文件
for label, kind in enumerate(["Type1", "Type2", "Type3", "Type4"]):
for path in glob(os.path.join(DATA_ROOT_PATH, "Type1/*.jpg")):
dataset.append({
'kind': kind,
'image_name': path.split('/')[-1],
'label': label
})
# 将数据顺序打乱
random.shuffle(dataset)
dataset = pd.DataFrame(dataset)
# K折交叉验证
gkf = GroupKFold(n_splits=5)
dataset.loc[:, 'fold'] = 0
for fold_number, (train_index, val_index) in enumerate(gkf.split(X=dataset.index, y=dataset['label'], groups=dataset['image_name'])):
dataset.loc[dataset.iloc[val_index].index, 'fold'] = fold_number
# 训练数据增强
def get_train_transforms():
return A.Compose([
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.5),
A.Resize(height=512, width=512, p=1.0),
ToTensorV2(p=1.0),
], p=1.0)
def get_valid_transforms():
return A.Compose([
A.Resize(height=512, width=512, p=1.0),
ToTensorV2(p=1.0),
], p=1.0)
# 独热处理
def onehot(size, target):
vec = torch.zeros(size, dtype=torch.float32)
vec[target] = 1.
return vec
class DatasetRetriever(Dataset):
def __init__(self, kinds, image_names, labels, transforms=None):
super().__init__()
self.kinds = kinds
self.image_names = image_names
self.labels = labels
self.transforms = transforms
def __getitem__(self, index: int):
kind, image_name, label = self.kinds[index], self.image_names[index], self.labels[index]
image = cv2.imread(f'{DATA_ROOT_PATH}/{kind}/{image_name}', cv2.IMREAD_COLOR)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)
image /= 255.0
if self.transforms:
sample = {'image': image}
sample = self.transforms(**sample)
image = sample['image']
target = onehot(4, label)
return image, target
def __len__(self) -> int:
return self.image_names.shape[0]
def get_labels(self):
return list(self.labels)
fold_number = 0
def get_dataset():
train_dataset = DatasetRetriever(
kinds=dataset[dataset['fold'] != fold_number].kind.values,
image_names=dataset[dataset['fold'] != fold_number].image_name.values,
labels=dataset[dataset['fold'] != fold_number].label.values,
transforms=get_train_transforms(),
)
validation_dataset = DatasetRetriever(
kinds=dataset[dataset['fold'] == fold_number].kind.values,
image_names=dataset[dataset['fold'] == fold_number].image_name.values,
labels=dataset[dataset['fold'] == fold_number].label.values,
transforms=get_valid_transforms(),
)
return train_dataset, validation_dataset
此处的修改内容为
DATA_ROOT_PATH = "data" # 数据集文件,修改成自己的数据集文件路径
for label, kind in enumerate(["自定义1", "自定义2", "自定义3", "自定义4"]):
# 改成自己分类文件名,数量可增加
for path in glob(os.path.join(DATA_ROOT_PATH, "自定义1/*.jpg")):
dataset.append({
'kind': kind,
'image_name': path.split('/')[-1],
'label': label
})
来看第二个代码块“utils.py”
from glob import glob
import torch
from torch import nn
import os
from datetime import datetime
import time
import albumentations as A
from albumentations.pytorch.transforms import ToTensorV2
import warnings
warnings.filterwarnings("ignore")
# 获取训练数据及训练数据增强
def get_train_transforms():
return A.Compose([
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.5),
A.Resize(height=512, width=512, p=1.0),
ToTensorV2(p=1.0),
], p=1.0)
# 测试集
def get_valid_transforms():
return A.Compose([
A.Resize(height=512, width=512, p=1.0),
ToTensorV2(p=1.0),
], p=1.0)
class AverageMeter(object):
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
# 标签平滑处理
class LabelSmoothing(nn.Module):
def __init__(self, smoothing=0.05):
super(LabelSmoothing, self).__init__()
self.confidence = 1.0 - smoothing
self.smoothing = smoothing
def forward(self, x, target):
if self.training:
x = x.float()
target = target.float()
logprobs = torch.nn.functional.log_softmax(x, dim=-1)
nll_loss = -logprobs * target
nll_loss = nll_loss.sum(-1)
smooth_loss = -logprobs.mean(dim=-1)
loss = self.confidence * nll_loss + self.smoothing * smooth_loss
return loss.mean()
else:
return torch.nn.functional.cross_entropy(x, target)
# 训练方法集成
class Fitter:
def __init__(self, model, device, config):
self.config = config
self.epoch = 0
self.base_dir = 'base'
self.log_path = f'{self.base_dir}/log.txt'
self.best_summary_loss = 10 ** 5
self.model = model
self.device = device
param_optimizer = list(self.model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
{'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.001},
{'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
]
self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=config.lr)
self.scheduler = config.SchedulerClass(self.optimizer, **config.scheduler_params)
self.criterion = LabelSmoothing().to(self.device)
self.log(f'Fitter prepared. Device is {self.device}')
def fit(self, train_loader, validation_loader):
for e in range(self.config.n_epochs):
if self.config.verbose:
lr = self.optimizer.param_groups[0]['lr']
timestamp = datetime.utcnow().isoformat()
self.log(f'\n{timestamp}\nLR: {lr}')
t = time.time()
summary_loss = self.train_one_epoch(train_loader)
self.log(
f'[RESULT]: Train. Epoch: {self.epoch}, summary_loss: {summary_loss.avg:.5f}, time: {(time.time() - t):.5f}')
self.save(f'{self.base_dir}/last-checkpoint.bin')
t = time.time()
summary_loss, final_scores = self.validation(validation_loader)
self.log(
f'[RESULT]: Val. Epoch: {self.epoch}, summary_loss: {summary_loss.avg:.5f}, time: {(time.time() - t):.5f}')
if summary_loss.avg < self.best_summary_loss:
self.best_summary_loss = summary_loss.avg
self.model.eval()
self.save(f'{self.base_dir}/best-checkpoint-{str(self.epoch).zfill(3)}epoch.bin')
for path in sorted(glob(f'{self.base_dir}/best-checkpoint-*epoch.bin'))[:-3]:
os.remove(path)
if self.config.validation_scheduler:
self.scheduler.step(metrics=summary_loss.avg)
self.epoch += 1
def validation(self, val_loader):
self.model.eval()
summary_loss = AverageMeter()
t = time.time()
for step, (images, targets) in enumerate(val_loader):
if self.config.verbose:
if step % self.config.verbose_step == 0:
print(
f'Val Step {step}/{len(val_loader)}, ' + \
f'summary_loss: {summary_loss.avg:.5f},' + \
f'time: {(time.time() - t):.5f}', end='\r'
)
with torch.no_grad():
targets = targets.to(self.device).float()
batch_size = images.shape[0]
images = images.to(self.device).float()
outputs = self.model(images)
loss = self.criterion(outputs, targets)
summary_loss.update(loss.detach().item(), batch_size)
return summary_loss
def train_one_epoch(self, train_loader):
self.model.train()
summary_loss = AverageMeter()
t = time.time()
for step, (images, targets) in enumerate(train_loader):
if self.config.verbose:
if step % self.config.verbose_step == 0:
print(
f'Train Step {step}/{len(train_loader)}, ' + \
f'summary_loss: {summary_loss.avg:.5f},' + \
f'time: {(time.time() - t):.5f}', end='\r')
targets = targets.to(self.device).float()
images = images.to(self.device).float()
batch_size = images.shape[0]
self.optimizer.zero_grad()
outputs = self.model(images)
loss = self.criterion(outputs, targets)
loss.backward()
summary_loss.update(loss.detach().item(), batch_size)
self.optimizer.step()
if self.config.step_scheduler:
self.scheduler.step()
return summary_loss,
def save(self, path):
self.model.eval()
torch.save({
'model_state_dict': self.model.state_dict(),
'optimizer_state_dict': self.optimizer.state_dict(),
'scheduler_state_dict': self.scheduler.state_dict(),
'best_summary_loss': self.best_summary_loss,
'epoch': self.epoch,
}, path)
def load(self, path):
checkpoint = torch.load(path)
self.model.load_state_dict(checkpoint['model_state_dict'])
self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
self.best_summary_loss = checkpoint['best_summary_loss']
self.epoch = checkpoint['epoch'] + 1
def log(self, message):
if self.config.verbose:
print(message)
with open(self.log_path, 'a+') as logger:
logger.write(f'{message}\n')
此处需要修改的内容
class Fitter:
def __init__(self, model, device, config):
self.config = config
self.epoch = 0
self.base_dir = 'base' # 此处需要保存模型和日志文件,选择自己用来保存模型和日志文件的文件夹路径
self.log_path = f'{self.base_dir}/log.txt'
self.best_summary_loss = 10 ** 5
self.model = model
self.device = device
再来看第三个代码块“train.py”
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.sampler import SequentialSampler, RandomSampler
from efficientnet_pytorch import EfficientNet
from catalyst.data.sampler import BalanceClassSampler
from utils import Fitter
from dataset import get_dataset
# 激活网络
def get_net():
net = EfficientNet.from_pretrained('efficientnet-b2') # 确定使用的网络,efficient-b0~b7
net._fc = nn.Linear(in_features=1408, out_features=4, bias=True) # out_features改为你所需要的分类数目
return net
net = get_net().cuda()
# 超参的定义
class TrainGlobalConfig:
num_workers = 4
batch_size = 8 # 显存够大的话可以适当增大batch_size
n_epochs = 25 # 训练的epochs
lr = 0.001 # 学习率的初始值
verbose = True
verbose_step = 1
step_scheduler = False
validation_scheduler = True
SchedulerClass = torch.optim.lr_scheduler.ReduceLROnPlateau
scheduler_params = dict(
mode='min',
factor=0.5,
patience=1,
verbose=False,
threshold=0.0001,
threshold_mode='abs',
cooldown=0,
min_lr=1e-8,
eps=1e-08
)
train_dataset, validation_dataset = get_dataset()
def run_training():
device = torch.device('cuda:0')
train_loader = torch.utils.data.DataLoader(
train_dataset,
sampler=BalanceClassSampler(labels=train_dataset.get_labels(), mode="downsampling"),
batch_size=TrainGlobalConfig.batch_size,
pin_memory=False,
drop_last=True,
num_workers=TrainGlobalConfig.num_workers,
)
val_loader = torch.utils.data.DataLoader(
validation_dataset,
batch_size=TrainGlobalConfig.batch_size,
num_workers=TrainGlobalConfig.num_workers,
shuffle=False,
sampler=SequentialSampler(validation_dataset),
pin_memory=False,
)
fitter = Fitter(model=net, device=device, config=TrainGlobalConfig)
# fitter.load(f'{fitter.base_dir}/last-checkpoint.bin')
fitter.fit(train_loader, val_loader)
run_training()
需要修改的部分
# 激活网络
def get_net():
net = EfficientNet.from_pretrained('efficientnet-b2') # 确定使用的网络,efficient-b0~b7
net._fc = nn.Linear(in_features=1408, out_features=4, bias=True) # out_features改为你所需要的分类数目
return net
net = get_net().cuda()
# 超参的定义
class TrainGlobalConfig:
num_workers = 4
batch_size = 8 # 显存够大的话可以适当增大batch_size
n_epochs = 25 # 训练的epochs
lr = 0.001 # 学习率的初始值
按照顺序依次修改上述代码,运行train.py文件,即可开始愉快的训练了
如果想看模型训练的效果,使用inference.py模块进行推断
import cv2
import torch
from torch import nn
from efficientnet_pytorch import EfficientNet
path = "图片路径"
def get_net():
net = EfficientNet.from_pretrained('efficientnet-b2')
net._fc = nn.Linear(in_features=1408, out_features=4, bias=True)
return net
images = cv2.imread(path)
net = get_net().cuda()
checkpoint = torch.load('base/best-checkpoint-024epoch.bin') # 训练得到的模型路径
net.load_state_dict(checkpoint['model_state_dict'])
net.eval()
y_pred = net(images.cuda())
y_pred = nn.functional.softmax(y_pred, dim=1).data.cpu().numpy() # 得到属于各个类别的概率
来看一下项目文件结构