基于ECANet18，迁移学习训练RAF_DB数据集

目录

• 写在前面
• 1、导入相关模块
• 2、显卡加速
• 3、设置相关数据
• 4、模型定义
• 5、 重写全连接层
• 6、导入模型权重
• 7、加载数据集
• 8、优化器，损失函数
• 9、测试函数
• 10、训练，加入AMP优化

代码相关技术

1、空间注意力
2、显卡加速
3、AMP加速
4、余弦退火动态学习率
5、迁移学习

写在前面

代码基于notebook写的，每个单元格是有顺序的，官方模型需要自己下载，
官方代码，权重下载：https://github.com/BangguWu/ECANet
个人感觉这个ECA模块很轻量，而且高效。
如果你有更好的模块架构或者炼丹技巧，你可以在评论区留言，我们一起优化代码。

1、导入相关模块

# 导包
import matplotlib.pyplot as plt
import math
from torch.cuda.amp import autocast 
from torch.cuda.amp import GradScaler
import torch
import os
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torch import optim
import torch.nn.functional as F
import numpy as np
import pandas as pd
import torch.nn as nn
import torchvision
from torchsummary import summary
import time
from torch.nn import init
from typing import Union, List, Dict, Any, Optional, cast
import warnings
warnings.filterwarnings("ignore")

2、显卡加速

torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True

3、设置相关数据

#保存训练数据和模型
data_csv_path = "./qyxx-eca_net.csv"   #修改此处文件名 
model_save_path = "./qyxx-eca_net.pkl"  #修改此处文件名
train_path = "./RAF-DB/train"
val_path = "./RAF-DB/test"
#模型批次大小
batch_size = 128
resume = True
#动态学习率，学习率和循环次数增加
lr = 5e-4
epochs = 200
D_epoch = 0 
best_acc  = 0
print("epochs:",epochs,"learning_rate:",lr,"batch_size:",batch_size)
#设备选取
flag = torch.cuda.is_available()
if flag:print("GPU")
else:print("CPU")
ngpu = 1
device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
# device = torch.device("cpu")
#查看显卡名称
#torch.cuda.get_device_name()
print("divice is ", device)#数据预处理（建议提前resize，减少每次资源的损失）放大到112x112 ，随机水平翻转
train_transform = transforms.Compose([transforms.Resize((224,224)),transforms.RandomHorizontalFlip(p=0.5),transforms.ToTensor(),transforms.Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225])
])#增加不同种transform，预测集中去除随机翻转
val_transform = transforms.Compose([transforms.Resize((224,224)),transforms.ToTensor(),transforms.Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225])
])

4、模型定义

class eca_layer(nn.Module):"""Constructs a ECA module.Args:channel: Number of channels of the input feature mapk_size: Adaptive selection of kernel size"""def __init__(self, channel, k_size=3):super(eca_layer, self).__init__()self.avg_pool = nn.AdaptiveAvgPool2d(1)self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False) self.sigmoid = nn.Sigmoid()def forward(self, x):# feature descriptor on the global spatial informationy = self.avg_pool(x)# Two different branches of ECA moduley = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)# Multi-scale information fusiony = self.sigmoid(y)return x * y.expand_as(x)
def conv3x3(in_planes, out_planes, stride=1):"""3x3 convolution with padding"""return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,padding=1, bias=False)class ECABasicBlock(nn.Module):expansion = 1def __init__(self, inplanes, planes, stride=1, downsample=None, k_size=3):super(ECABasicBlock, self).__init__()self.conv1 = conv3x3(inplanes, planes, stride)self.bn1 = nn.BatchNorm2d(planes)self.relu = nn.ReLU(inplace=True)self.conv2 = conv3x3(planes, planes, 1)self.bn2 = nn.BatchNorm2d(planes)self.eca = eca_layer(planes, k_size)self.downsample = downsampleself.stride = stridedef forward(self, x):residual = xout = self.conv1(x)out = self.bn1(out)out = self.relu(out)out = self.conv2(out)out = self.bn2(out)out = self.eca(out)if self.downsample is not None:residual = self.downsample(x)out += residualout = self.relu(out)return outclass ECABottleneck(nn.Module):expansion = 4def __init__(self, inplanes, planes, stride=1, downsample=None, k_size=3):super(ECABottleneck, self).__init__()self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)self.bn1 = nn.BatchNorm2d(planes)self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,padding=1, bias=False)self.bn2 = nn.BatchNorm2d(planes)self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)self.bn3 = nn.BatchNorm2d(planes * 4)self.relu = nn.ReLU(inplace=True)self.eca = eca_layer(planes * 4, k_size)self.downsample = downsampleself.stride = stridedef forward(self, x):residual = xout = self.conv1(x)out = self.bn1(out)out = self.relu(out)out = self.conv2(out)out = self.bn2(out)out = self.relu(out)out = self.conv3(out)out = self.bn3(out)out = self.eca(out)if self.downsample is not None:residual = self.downsample(x)out += residualout = self.relu(out)return outclass ResNet(nn.Module):def __init__(self, block, layers, num_classes=1000, k_size=[3, 3, 3, 3]):self.inplanes = 64super(ResNet, self).__init__()self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,bias=False)self.bn1 = nn.BatchNorm2d(64)self.relu = nn.ReLU(inplace=True)self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)self.layer1 = self._make_layer(block, 64, layers[0], int(k_size[0]))self.layer2 = self._make_layer(block, 128, layers[1], int(k_size[1]), stride=2)self.layer3 = self._make_layer(block, 256, layers[2], int(k_size[2]), stride=2)self.layer4 = self._make_layer(block, 512, layers[3], int(k_size[3]), stride=2)self.avgpool = nn.AvgPool2d(7, stride=1)self.fc = nn.Linear(512 * block.expansion, num_classes)for m in self.modules():if isinstance(m, nn.Conv2d):n = m.kernel_size[0] * m.kernel_size[1] * m.out_channelsm.weight.data.normal_(0, math.sqrt(2. / n))elif isinstance(m, nn.BatchNorm2d):m.weight.data.fill_(1)m.bias.data.zero_()def _make_layer(self, block, planes, blocks, k_size, stride=1):downsample = Noneif stride != 1 or self.inplanes != planes * block.expansion:downsample = nn.Sequential(nn.Conv2d(self.inplanes, planes * block.expansion,kernel_size=1, stride=stride, bias=False),nn.BatchNorm2d(planes * block.expansion),)layers = []layers.append(block(self.inplanes, planes, stride, downsample, k_size))self.inplanes = planes * block.expansionfor i in range(1, blocks):layers.append(block(self.inplanes, planes, k_size=k_size))return nn.Sequential(*layers)def forward(self, x):x = self.conv1(x)x = self.bn1(x)x = self.relu(x)x = self.maxpool(x)x = self.layer1(x)x = self.layer2(x)x = self.layer3(x)x = self.layer4(x)x = self.avgpool(x)x = x.view(x.size(0), -1)x = self.fc(x)return xdef eca_resnet18(k_size=[3, 5, 7, 7], num_classes=1_000, pretrained=False):"""Constructs a ResNet-18 model.Args:k_size: Adaptive selection of kernel sizepretrained (bool): If True, returns a model pre-trained on ImageNetnum_classes:The classes of classification"""model = ResNet(ECABasicBlock, [2, 2, 2, 2], num_classes=num_classes, k_size=k_size)model.avgpool = nn.AdaptiveAvgPool2d(1)return model

5、 重写全连接层

class SKNet(nn.Module):def __init__(self, num_class=7):super(SKNet, self).__init__()self.features = nn.Sequential(*list(eca_net.children())[:-2])self.avgpool = nn.AvgPool2d(7, stride=1)self.fc = nn.Linear(512, num_class)def forward(self, x):x = self.features(x)out = self.avgpool(x)out = torch.flatten(out,1)out = self.fc(out)return out

6、导入模型权重

from torch.nn.parallel import DistributedDataParallel
model_path = "C:\\Users\\sk\\Desktop\\eca_resnet18_k3577.pth.tar"
eca_net = eca_resnet18()
checkpoint = torch.load(model_path)
eca_net.load_state_dict({k.replace('module.',''):v for k,v in checkpoint['state_dict'].items()})
del checkpoint
eca_net

7、加载数据集

#使用torchvision.datasets.ImageFolder读取数据集指定train和test文件夹
train_data = torchvision.datasets.ImageFolder(train_path, transform=train_transform)
#drop_last舍弃未满一个批次的数据        num_workers工作区一般设置为GPU个数的4倍
data0_train = DataLoader(train_data, batch_size=batch_size, shuffle=True, drop_last=True,num_workers=4)
print(train_data)  #输出训练集相关
val_data = torchvision.datasets.ImageFolder(val_path, transform=val_transform)
data1_val = DataLoader(val_data, batch_size=batch_size, shuffle=True,drop_last=True,num_workers=4)
print(val_data)  #输出测试集相关print(train_data.classes)  #根据分的文件夹的名字来确定的类别
print(train_data.class_to_idx) #按顺序为这些类别定义索引为0,1...
print()
print(val_data.classes)
print(val_data.class_to_idx)
print()

8、优化器，损失函数

#to(device)将模型加入GPU中加速计算
model = SKNet().to(device)
#设置优化器
params = list(model.parameters()) 
optimizer = optim.AdamW(params, lr=lr)
# 如果不用SGD的话，Adam是最好的选择，学习率5e-4一般是比较优的
# optimizer = optim.Adam(model.parameters(), lr=lr)
#设置损失函数
criteon = nn.CrossEntropyLoss().to(device)
#余弦衰减学习率 
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=80, eta_min=0)
#形如TensorFlow中的summary函数输出模型参数
summary(model, input_size=[(3, 224, 224)], batch_size=batch_size, device="cuda")
print()

9、测试函数

#测试函数
def evalute_(model,val_loader):model.eval()test_loss2 = 0.0test_corrects2 = 0.0number = 0for batchidx, (x, label) in enumerate(val_loader):
#         print(number)#torch.cuda.empty_cache()  #清除非必要GPU缓存，但是我建议不要在训练中使用此句，这可能会损失你相当多的时间number = number + 1x, label = x.to(device), label.to(device)#测试函数中加入no_grad()，如果不加会增加计算和显存with torch.no_grad():y1 = model(x)#虽然可以直接使用max函数，但是我建议在y1的比较重你最好使用F.softmax(y1,dim=1)，这样可能会有更好的效果，我在训练中使用了它_, preds1 = torch.max(F.softmax(y1,dim=1), 1)loss = criteon(y1, label)  test_loss2 += loss.item()*batch_sizetest_corrects2 += torch.sum(preds1 == label.data)#由于使用了最后一次抛弃，我不能使用全部测试集作为分母，这样会使最后的准确率变小test_loss1 = test_loss2 / (number*batch_size)test_acc1 = test_corrects2.double() / (number*batch_size)
#     print("TestDataset loss is ", test_loss1,"TestDataset accuracy is ",test_acc1)return test_acc1, test_loss1
print("执行结束")

10、训练，加入AMP优化

#关于AMP自动精度求解，我也并不是很熟悉，只能使用官方给的实例进行照葫芦画瓢。
scaler = torch.cuda.amp.GradScaler()
torch.cuda.empty_cache()
for epoch in range(D_epoch, epochs):time_one = time.time()                         #标记训练开始时间戳train_acc1 = 0.0train_loss1 = 0.0train_acc = 0.0train_loss = 0.0val_acc = 0.0val_loss = 0.0number = 0model.train()print("epoch:",epoch)for batchidx , (x ,label) in enumerate(data0_train):x , label = x.to(device), label.to(device)optimizer.zero_grad()with torch.cuda.amp.autocast():y1 = model(x)loss = criteon(y1,label) _, preds1 = torch.max(F.softmax(y1,dim=1), 1)#AMP优化scaler.scale(loss).backward()
#         loss.backward()#梯度剪裁 有时候加了比没加还差，emmmm，你懂得torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm = 5, norm_type=2)  scaler.step(optimizer)
#         optimizer.step()scaler.update()train_loss1 += loss.item()*batch_sizetrain_acc1 += torch.sum(preds1 == label.data).double()number = number + 1time_two = time.time()             #标记训练结束时间戳#输出训练一轮所需时间，用于分析对比print("The Model-Train-Time spent  %d min %.2f s"%((time_two-time_one)//60,(time_two-time_one)%60))#计算训练时候的平均损失和平均准确率train_loss = train_loss1 / (number*batch_size)train_acc = train_acc1 / (number*batch_size)#计算测试时候的平均损失和平均准确率val_acc, val_loss = evalute_(model, data1_val)train_acc = train_acc.cpu()val_acc = val_acc.cpu()print('Accuracy : Train is {} , Valid is {} ;  Loss : Train is  {} ,Valid is {}'.format(train_acc, val_acc, train_loss , val_loss))#如果你不需要训练以及验证的准确率和损失值，你可以注释这下面的两行，它们不是非必须的，理论上只存在于汇报和论文中dataframe = pd.DataFrame(columns = [epoch,train_acc,train_loss,val_acc, val_loss])dataframe.to_csv(data_csv_path,line_terminator="\n",mode='a',index=False,sep=',')if val_acc > best_acc:print("覆盖最好的模型...")best_acc = val_acc checkpoint = {'epoch': epoch,'model': model.state_dict(),'optimizer': optimizer.state_dict(),'best_acc':best_acc}torch.save(checkpoint,model_save_path)
#     time_three = time.time() 
#     print("测试花费时间",time_three-time_two)scheduler.step()  #动态学习率更新 
#如果你不是非必须，我建议你尽量不要使用n折交叉验证，使用数据增强可能效果更优于它。

本文来自互联网用户投稿，文章观点仅代表作者本人，不代表本站立场，不承担相关法律责任。如若转载，请注明出处。 如若内容造成侵权/违法违规/事实不符，请点击【内容举报】进行投诉反馈！