@

导入项目使用的库

设置全局参数

图像预处理与增强

读取数据

设置模型

定义训练和验证函数

训练函数

验证函数

调用训练和验证方法

再次训练的模型为什么只保存model.state_dict()

在上一篇文章中完成了前期的准备工作，见链接：RepGhost实战：使用RepGhost实现图像分类任务(一）这篇主要是讲解如何训练和测试

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

import jsonimport osimport matplotlib.pyplot as pltimport torchimport torch.nn as nnimport torch.nn.parallelimport torch.optim as optimimport torch.utils.dataimport torch.utils.data.distributedimport torchvision.transforms as transformsfrom timm.utils import accuracy, AverageMeter, ModelEmafrom sklearn.metrics import classification_reportfrom timm.data.mixup import Mixupfrom timm.loss import SoftTargetCrossEntropyfrom torch.autograd import Variablefrom torchvision import datasetsfrom model import repghostnet_1_0xtorch.backends.cudnn.benchmark = Falseimport warningswarnings.filterwarnings("ignore")

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

if __name__ == '__main__':    #创建保存模型的文件夹    file_dir = 'checkpoints/RepGhost'    if os.path.exists(file_dir):        print('true')

        os.makedirs(file_dir,exist_ok=True)    else:        os.makedirs(file_dir)

    # 设置全局参数    model_lr = 1e-4    BATCH_SIZE = 16    EPOCHS = 1000    DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')    use_amp = True  # 是否使用混合精度    use_dp = False #是否开启dp方式的多卡训练    classes = 12    resume =None    CLIP_GRAD = 5.0    Best_ACC = 0 #记录最高得分    use_ema=True    model_ema_decay=0.9998    start_epoch=1

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

接下来，查看全局参数：

model_lr：学习率，根据实际情况做调整。

BATCH_SIZE：batchsize，根据显卡的大小设置。

EPOCHS：epoch的个数，一般300够用。

use_amp：是否使用混合精度。

classes：类别个数。

resume：是否接着上次模型继续训练，如果不为空，则按照resume的值加载模型。如果是None则表示不接着上次训练的模型训练。

CLIP_GRAD：梯度的最大范数，在梯度裁剪里设置。

Best_ACC：记录最高ACC得分。 use_ema：是否使用ema model_ema_decay：ema衰减值， start_epoch：开始训练的epoch，默认是1，如果再次训练的时候，加载模型的时候会自动对其设置。

数据处理比较简单，加入了随机10度的旋转、高斯模糊、颜色的调整、做了Resize和归一化，定义Mixup函数。

这里注意下Resize的大小，由于选用的Ghost模型输入是224×224的大小，所以要Resize为224×224。

 # 数据预处理7    transform = transforms.Compose([        transforms.RandomRotation(10),        transforms.GaussianBlur(kernel_size=(5,5),sigma=(0.1, 3.0)),        transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5),        transforms.Resize((224, 224)),        transforms.ToTensor(),        transforms.Normalize(mean=[0.44127703, 0.4712498, 0.43714803], std= [0.18507297, 0.18050247, 0.16784933])

    ])    transform_test = transforms.Compose([        transforms.Resize((224, 224)),        transforms.ToTensor(),        transforms.Normalize(mean=[0.44127703, 0.4712498, 0.43714803], std= [0.18507297, 0.18050247, 0.16784933])    ])    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)    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))    # 导入数据    train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True,drop_last=True)    test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=False)

在train_loader中将drop_last设置为True了，舍弃最后一次迭代，因为MixUp要求输入的Batchsize是偶数，但是最后一次迭代有可能是奇数所以就会报错。 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()。

设置模型为ghostnet_100，num_classes设置为12。如果resume为True，则加载模型接着上次训练。在val阶段保存的模型，在这里要注意，我只保存了模型的权重参数，没有保存整个模型。Best模型只保存了权重参数，在保存每个epoch的模型时，除了保存权重参数外，还保存了 Best_ACC和epoch信息，方便再次训练。

如果resume不为None，则表示再次载入模型训练，将模型的权重参数载入，再次训练。

优化器设置为adamW。

学习率调整策略选择为余弦退火。

开启混合精度训练,声明pytorch自带的混合精度 torch.cuda.amp.GradScaler()。

检测可用显卡的数量，如果大于1，并且开启多卡训练的情况下，则要用torch.nn.DataParallel加载模型，开启多卡训练。

use_ema，是否使用ema，如果为True，则初始化ema。

     # 实例化模型并且移动到GPU    criterion_train = SoftTargetCrossEntropy()    criterion_val = torch.nn.CrossEntropyLoss()    #设置模型    model_ft = repghostnet_1_0x()    print(model_ft)    num_ftrs = model_ft.classifier.in_features    model_ft.classifier = nn.Linear(num_ftrs, classes)    if resume:        model=torch.load(resume)        print(model['state_dict'].keys())        model_ft.load_state_dict(model['state_dict'])        Best_ACC=model['Best_ACC']        start_epoch=model['epoch']+1    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:        scaler = torch.cuda.amp.GradScaler()    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)    if use_ema:        model_ema = ModelEma(            model_ft,            decay=model_ema_decay,            device=DEVICE,            resume=resume)    else:        model_ema=None

注：torch.nn.DataParallel方式，默认不能开启混合精度训练的，如果想要开启混合精度训练，则需要在模型的forward前面加上@autocast()函数。导入包from torch.cuda.amp import autocast，如果是cpu，则导入from torch.cpu.amp import autocast

如果不开启混合精度则要将@autocast()去掉，否则loss一直试nan。如果不开DP模式也要去掉。

训练的主要步骤：

1、使用AverageMeter保存自定义变量，包括loss，ACC1，ACC5。

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

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

4、如果使用混合精度，则

with torch.cuda.amp.autocast(),开启混合精度。

计算loss。

scaler.scale(loss).backward()，梯度放大。

torch.nn.utils.clip_grad_norm_，梯度裁剪，放置梯度爆炸。

scaler.step(optimizer) ，首先把梯度值unscale回来，如果梯度值不是inf或NaN,则调用optimizer.step()来更新权重，否则，忽略step调用，从而保证权重不更新。

更新下一次迭代的scaler。

否则，直接反向传播求梯度。torch.nn.utils.clip_grad_norm_函数执行梯度裁剪，防止梯度爆炸。

5、如果model_ema不为None，则执行model_ema的updata函数，更新模型。

6、 torch.cuda.synchronize()，等待上面所有的操作执行完成。

7、接下来，更新loss，ACC1，ACC5的值。

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

# 定义训练过程def train(model, device, train_loader, optimizer, epoch,model_ema):    model.train()    loss_meter = AverageMeter()    acc1_meter = AverageMeter()    acc5_meter = AverageMeter()    total_num = len(train_loader.dataset)    print(total_num, len(train_loader))    for batch_idx, (data, target) in enumerate(train_loader):        data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)        samples, targets = mixup_fn(data, target)        output = model(samples)        optimizer.zero_grad()        if use_amp:            with torch.cuda.amp.autocast():                loss = criterion_train(output, targets)            scaler.scale(loss).backward()            torch.nn.utils.clip_grad_norm_(model.parameters(), CLIP_GRAD)            # Unscales gradients and calls            # or skips optimizer.step()            scaler.step(optimizer)            # Updates the scale for next iteration            scaler.update()        else:            loss = criterion_train(output, targets)            loss.backward()            # torch.nn.utils.clip_grad_norm_(model.parameters(), CLIP_GRAD)            optimizer.step()        if model_ema is not None:            model_ema.update(model)        torch.cuda.synchronize()        lr = optimizer.state_dict()['param_groups'][0]['lr']        loss_meter.update(loss.item(), target.size(0))        acc1, acc5 = accuracy(output, target, topk=(1, 5))        loss_meter.update(loss.item(), target.size(0))        acc1_meter.update(acc1.item(), target.size(0))        acc5_meter.update(acc5.item(), target.size(0))        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 =loss_meter.avg    acc = acc1_meter.avg    print('epoch:{}\tloss:{:.2f}\tacc:{:.2f}'.format(epoch, ave_loss, acc))    return ave_loss, acc

验证集和训练集大致相似，主要步骤：

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

2、在val的函数上面添加@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、调用accuracy函数计算ACC1和ACC5

5、更新loss_meter、acc1_meter、acc5_meter的参数。

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

如果acc比Best_ACC大，则保存模型。 保存模型的逻辑：

如果ACC比Best_ACC高，则保存best模型 判断模型是否为DP方式训练的模型。

如果是DP方式训练的模型，模型参数放在model.module，则需要保存model.module。 否则直接保存model。 这里注意：只保存了model的参数，没有整个模型，因为Rep模型还要经历转换，转换代码里面使用load_state_dict加载的，所以只能保存模型的参数。 接下来保存每个epoch的模型。 判断模型是否为DP方式训练的模型。 如果是DP方式训练的模型，模型参数放在model.module，则需要保存model.module.state_dict()。 新建个字典，放置Best_ACC、epoch和 model.module.state_dict()等参数。然后将这个字典保存。 否则，新建个字典，放置Best_ACC、epoch和 model.state_dict()等参数。然后将这个字典保存。 在这里注意：对于每个epoch的模型只保存了state_dict参数，没有保存整个模型文件。

# 验证过程@torch.no_grad()def val(model, device, test_loader):    global Best_ACC    model.eval()    loss_meter = AverageMeter()    acc1_meter = AverageMeter()    acc5_meter = AverageMeter()    total_num = len(test_loader.dataset)    print(total_num, len(test_loader))    val_list = []    pred_list = []

    for data, target in test_loader:        for t in target:            val_list.append(t.data.item())        data, target = data.to(device,non_blocking=True), target.to(device,non_blocking=True)        output = model(data)        loss = criterion_val(output, target)        _, pred = torch.max(output.data, 1)        for p in pred:            pred_list.append(p.data.item())        acc1, acc5 = accuracy(output, target, topk=(1, 5))        loss_meter.update(loss.item(), target.size(0))        acc1_meter.update(acc1.item(), target.size(0))        acc5_meter.update(acc5.item(), target.size(0))    acc = acc1_meter.avg    print('\nVal set: Average loss: {:.4f}\tAcc1:{:.3f}%\tAcc5:{:.3f}%\n'.format(        loss_meter.avg,  acc,  acc5_meter.avg))

    if acc > Best_ACC:        if isinstance(model, torch.nn.DataParallel):            torch.save(model.module.state_dict(), file_dir + '/' + 'best.pth')        else:            torch.save(model.state_dict(), file_dir + '/' + 'best.pth')        Best_ACC = acc    if isinstance(model, torch.nn.DataParallel):        state = {

            'epoch': epoch,            'state_dict': model.module.state_dict(),            'Best_ACC':Best_ACC        }        torch.save(state, file_dir + "/" + 'model_' + str(epoch) + '_' + str(round(acc, 3)) + '.pth')    else:        state = {            'epoch': epoch,            'state_dict': model.state_dict(),            'Best_ACC': Best_ACC        }        torch.save(state, file_dir + "/" + 'model_' + str(epoch) + '_' + str(round(acc, 3)) + '.pth')    return val_list, pred_list, loss_meter.avg, 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、如果use_ema为True，则使用ema，所以将ema模型传入验证函数，否则将model_ft模型传入验证函数。 3、调用验证函数，得到val_list, pred_list, val_loss, val_acc。将val_loss, val_acc分别放入val_loss_list和val_acc_list中，然后存入到logdir字典中。

4、保存log。

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

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

7、绘制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,model_ema)        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        if use_ema:            val_list, pred_list, val_loss, val_acc = val(model_ema.ema, DEVICE, test_loader)        else:            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