> 文档中心 > RepLKNet实战:使用RepLKNet实现对植物幼苗的分类(非官方)(二)

RepLKNet实战:使用RepLKNet实现对植物幼苗的分类(非官方)(二)


训练

完成上面的步骤后,就开始train脚本的编写,新建train.py.

导入项目使用的库

import jsonimport osimport shutilimport matplotlib.pyplot as pltimport torchimport torch.nn.parallelimport torch.optim as optimimport torch.utils.dataimport torch.utils.data.distributedimport torchvision.transforms as transformsfrom apex import ampfrom sklearn.metrics import classification_reportfrom timm.data.mixup import Mixupfrom timm.loss import SoftTargetCrossEntropyfrom torchtoolbox.transform import Cutoutfrom torchvision import datasetsfrom models.replknet import create_RepLKNet31Btorch.backends.cudnn.benchmark = Falseimport warningswarnings.filterwarnings("ignore")os.environ['CUDA_VISIBLE_DEVICES'] = "0,1"

设置全局参数

设置学习率、BatchSize、epoch等参数,判断环境中是否存在GPU,如果没有则使用CPU。建议使用GPU,CPU太慢了。

if __name__ == '__main__':    #创建保存模型的文件夹    file_dir='checkpoints/replknet'    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)    # 设置全局参数    model_lr = 1e-4    BATCH_SIZE = 4    EPOCHS = 1000    DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')    use_amp = False  # 是否使用混合精度    use_dp=True #是否开启dp方式的多卡训练    classes = 12    resume = False    CLIP_GRAD = 5.0    model_path = 'best.pth'    Best_ACC = 0 #记录最高得分

设置存放权重文件的文件夹,如果文件夹存在删除再建立。

接下来,查看全局参数:

model_lr:学习率,根据实际情况做调整。

BATCH_SIZE:batchsize,根据显卡的大小设置。

EPOCHS:epoch的个数,一般300够用。

use_amp:是否使用混合精度。

classes:类别个数。

resume:是否接着上次模型继续训练。

model_path:模型的路径。如果resume设置为True时,就采用model_path定义的模型继续训练。

CLIP_GRAD:梯度的最大范数,在梯度裁剪里设置。

Best_ACC:记录最高ACC得分。

图像预处理与增强

数据处理比较简单,加入了Cutout、做了Resize和归一化,定义Mixup函数

这里注意下Resize的大小,由于RepLKNet的输入是224×224的大小,所以要Resize为224×224。

# 数据预处理7transform = transforms.Compose([    transforms.Resize((224, 224)),    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((224, 224)),    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文件中。

# 读取数据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()。
  • 设置模型为RepLKNet31B,然后加载预训练模型,num_classes设置为12。如果resume为True,则加载模型接着上次训练。
  • 优化器设置为adamW。
  • 学习率调整策略选择为余弦退火。
  • 开启混合精度训练。
  • 检测可用显卡的数量,如果大于1,并且开启多卡训练的情况下,则要用torch.nn.DataParallel加载模型,开启多卡训练。
  # 实例化模型并且移动到GPU    criterion_train = SoftTargetCrossEntropy()    criterion_val = torch.nn.CrossEntropyLoss()    # 设置模型    model_ft = create_RepLKNet31B()    model_ft.load_state_dict(torch.load('RepLKNet-31B_ImageNet-1K_224.pth'))    numftr=model_ft.head.in_features    model_ft.head=torch.nn.Linear(numftr,classes)    if resume: model_ft = torch.load(model_path)    model_ft.to(DEVICE)    print(model_ft)    # 选择简单暴力的Adam优化器,学习率调低    optimizer = optim.AdamW(model_ft.parameters(), lr=model_lr)    cosine_schedule = optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer, T_max=20, eta_min=1e-6)    if use_amp: model_ft, optimizer = amp.initialize(model_ft, optimizer, opt_level="O1")  # 这里是“欧一”,不是“零一”    if torch.cuda.device_count() > 1 and use_dp: print("Let's use", torch.cuda.device_count(), "GPUs!") model_ft = torch.nn.DataParallel(model_ft)

定义训练和验证函数

训练函数

训练的主要步骤:

1、判断迭代的数据是否是奇数,由于mixup_fn只能接受偶数,所以如果不是偶数则要减去一位,让其变成偶数。但是有可能最后一次迭代只有一条数据,减去后就变成了0,所以还要判断不能小于2,如果小于2则直接中断本次循环。

2、将数据输入mixup_fn生成mixup数据,然后输入model计算loss。

3、 optimizer.zero_grad() 梯度清零,把loss关于weight的导数变成0。

4、如果使用混合精度,则使用amp.scale_loss反向传播求解梯度,否则,直接反向传播求梯度。torch.nn.utils.clip_grad_norm_函数执行梯度裁剪,防止梯度爆炸。

5、 optimizer.step()更新参数。

6、接下来,获取学习率,获取loss、计算本次Batch的ACC

等待一个epoch训练完成后,计算平均loss和平均acc

# 定义训练过程def train(model, device, train_loader, optimizer, epoch):    model.train()    sum_loss = 0    correct = 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(samples) 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 _, pred = torch.max(output.data, 1) correct += torch.sum(pred == target) 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)    correct = correct.data.item()    acc = correct / total_num    print('epoch:{}\tloss:{}\tacc:{}'.format(epoch, ave_loss, acc))    return ave_loss, acc

验证函数

验证集和训练集大致相似,主要步骤:

1、定义参数,test_loss测试的loss,total_num总的验证集的数量,val_list验证集的label,pred_list预测的label。

2、在 with torch.no_grad()下面循环验证集,在该模块下,所有计算得出的tensor的requires_grad都自动设置为False。即使一个tensor(命名为x)的requires_grad = True,在with torch.no_grad计算,由x得到的新tensor(命名为w-标量)requires_grad也为False,且grad_fn也为None,即不会对w求导。

3、使用验证集的loss函数求出验证集的loss。

4、求出本次batch的acc

本次epoch循环完成后,求得本次epoch的acc、loss。

如果acc比Best_ACC大,则保存模型。

# 验证过程def val(model, device, test_loader):    global Best_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 > Best_ACC: if isinstance(model, torch.nn.DataParallel):     torch.save(model.module, file_dir + "/" + 'model_' + str(epoch) + '_' + str(round(acc, 3)) + '.pth')     torch.save(model.module, file_dir + '/' + 'best.pth') else:     torch.save(model, file_dir + "/" + 'model_' + str(epoch) + '_' + str(round(acc, 3)) + '.pth')     torch.save(model, file_dir + '/' + 'best.pth') Best_ACC = acc    return val_list, pred_list, avgloss, acc

调用训练和验证方法

调用训练函数和验证函数的主要步骤:

1、定义参数:

  • is_set_lr,是否已经设置了学习率,当epoch大于一定的次数后,会将学习率设置到一定的值,并将其置为True。
  • log_dir:记录log用的,将有用的信息保存到字典中,然后转为json保存起来。
  • train_loss_list:保存每个epoch的训练loss。
  • val_loss_list:保存每个epoch的验证loss。
  • train_acc_list:保存每个epoch的训练acc。
  • val_acc_list:保存么每个epoch的验证acc。
  • epoch_list:存放每个epoch的值。

循环epoch

1、调用train函数,得到 train_loss, train_acc,并将分别放入train_loss_list,train_acc_list,然后存入到logdir字典中。

2、调用验证函数,得到val_list, pred_list, val_loss, val_acc。将val_loss, val_acc分别放入val_loss_list和val_acc_list中,然后存入到logdir字典中。

3、保存log。

4、打印本次的测试报告。

5、如果epoch大于600,将学习率设置为固定的1e-6。

6、绘制loss曲线和acc曲线。

 # 训练与验证    is_set_lr = False    log_dir = {}    train_loss_list, val_loss_list, train_acc_list, val_acc_list, epoch_list = [], [], [], [], []    for epoch in range(1, EPOCHS + 1): epoch_list.append(epoch) train_loss, train_acc = train(model_ft, DEVICE, train_loader, optimizer, epoch) train_loss_list.append(train_loss) train_acc_list.append(train_acc) log_dir['train_acc'] = train_acc_list log_dir['train_loss'] = train_loss_list val_list, pred_list, val_loss, val_acc = val(model_ft, DEVICE, test_loader) val_loss_list.append(val_loss) val_acc_list.append(val_acc) log_dir['val_acc'] = val_acc_list log_dir['val_loss'] = val_loss_list log_dir['best_acc'] = Best_ACC with open(file_dir + '/result.json', 'w', encoding='utf-8') as file:     file.write(json.dumps(log_dir)) print(classification_report(val_list, pred_list, target_names=dataset_train.class_to_idx)) if epoch < 600:     cosine_schedule.step() else:     if not is_set_lr:  for param_group in optimizer.param_groups:      param_group["lr"] = 1e-6      is_set_lr = True fig = plt.figure(1) plt.plot(epoch_list, train_loss_list, 'r-', label=u'Train Loss') # 显示图例 plt.plot(epoch_list, val_loss_list, 'b-', label=u'Val Loss') plt.legend(["Train Loss", "Val Loss"], loc="upper right") plt.xlabel(u'epoch') plt.ylabel(u'loss') plt.title('Model Loss ') plt.savefig(file_dir + "/loss.png") plt.close(1) fig2 = plt.figure(2) plt.plot(epoch_list, train_acc_list, 'r-', label=u'Train Acc') plt.plot(epoch_list, val_acc_list, 'b-', label=u'Val Acc') plt.legend(["Train Acc", "Val Acc"], loc="lower right") plt.title("Model Acc") plt.ylabel("acc") plt.xlabel("epoch") plt.savefig(file_dir + "/acc.png") plt.close(2)

运行以及结果查看

完成上面的所有代码就可以开始运行了。点击右键,然后选择“run train.py”即可,运行结果如下:

RepLKNet实战:使用RepLKNet实现对植物幼苗的分类(非官方)(二)

在每个epoch测试完成之后,打印验证集的acc、recall等指标。

RepLKNet实战:使用RepLKNet实现对植物幼苗的分类(非官方)(二)
绘制acc曲线
RepLKNet实战:使用RepLKNet实现对植物幼苗的分类(非官方)(二)

绘制loss曲线
经过训练,最好的成绩能达到97%

测试

测试,我们采用一种通用的方式。

测试集存放的目录如下图:
RepLKNet实战:使用RepLKNet实现对植物幼苗的分类(非官方)(二)

import torch.utils.data.distributedimport torchvision.transforms as transformsfrom PIL import Imagefrom torch.autograd import Variableimport osclasses = ('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("checkpoints/replknet/best.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()]))

测试的主要逻辑:

1、定义类别,这个类别的顺序和训练时的类别顺序对应,一定不要改变顺序!!!!

2、定义transforms,transforms和验证集的transforms一样即可,别做数据增强。

3、 加载model,并将模型放在DEVICE里,

4、循环 读取图片并预测图片的类别,在这里注意,读取图片用PIL库的Image。不要用cv2,transforms不支持。循环里面的主要逻辑:

  • 使用Image.open读取图片
  • 使用transform_test对图片做归一化和标椎化。
  • img.unsqueeze_(0) 增加一个维度,由(3,224,224)变为(1,3,224,224)
  • Variable(img).to(DEVICE):将数据放入DEVICE中。
  • model(img):执行预测。
  • _, pred = torch.max(out.data, 1):获取预测值的最大下角标。

运行结果:

RepLKNet实战:使用RepLKNet实现对植物幼苗的分类(非官方)(二)

完整的代码

https://download.csdn.net/download/hhhhhhhhhhwwwwwwwwww/85456919?spm=1001.2014.3001.5503