参考链接ClassficationRetrainingAndInfer

训练代码解析

脚本名称：main.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import os
import time
import torch_npu
import torchvision.datasets as datasets
import torchvision.models as models
from torch_npu.npu import amp
from torch.utils.tensorboard import SummaryWriter
import datetime
import torchvision.transforms as transforms
import shutil

model_path = "models"
device = torch.device('npu:0')
tensorboard = SummaryWriter(log_dir=os.path.join(model_path, "tensorboard", f"{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}"))
best_accuracy = 0


class AverageMeter(object):
    """
    Computes and stores the average and current value
    """
    def __init__(self, name, fmt=':f'):
        self.name = name  # 名称
        self.fmt = fmt  # 格式化字符串
        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  # 计算平均值

    def __str__(self):
        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'  # 格式化字符串
        return fmtstr.format(**self.__dict__)  # 返回格式化后的字符串


class ProgressMeter(object):
    """
    Progress metering
    """
    def __init__(self, num_batches, meters, prefix=""):
        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)  # 初始化批次格式字符串
        self.meters = meters  # 计量器列表
        self.prefix = prefix  # 前缀字符串

    def display(self, batch):
        entries = [self.prefix + self.batch_fmtstr.format(batch)]  # 创建显示条目列表，首先添加带批次信息的前缀
        entries += [str(meter) for meter in self.meters]  # 将每个计量器的字符串表示添加到条目列表中
        print('  '.join(entries))  # 打印所有条目，用两个空格分隔

    def _get_batch_fmtstr(self, num_batches):
        num_digits = len(str(num_batches // 1))  # 计算批次数量的位数
        fmt = '{:' + str(num_digits) + 'd}'  # 创建格式化字符串，用于批次编号的对齐
        return '[' + fmt + '/' + fmt.format(num_batches) + ']'  # 返回格式化的批次字符串，例如 [  1/100]

def accuracy(output, target):
    """
    Computes the accuracy of predictions vs groundtruth
    """
    with torch.no_grad():  # 不计算梯度，以加速计算并减少内存使用

        output = F.softmax(output, dim=-1)  # 对模型输出应用softmax函数，获取概率分布
        _, preds = torch.max(output, dim=-1)  # 获取概率最高的类别的索引
        preds = (preds == target)  # 将预测结果与真实标签进行比较，得到一个布尔值数组
            
        return preds.float().mean().cpu().item() * 100.0  # 计算准确率：将布尔值转换为浮点数，计算平均值，转移到CPU，转换为Python数值，并乘以100转换为百分比
        
def train(train_loader, model, criterion, optimizer,scaler, epoch):
    """
    Train one epoch over the dataset
    """
    batch_time = AverageMeter('Time', ':6.3f')  # 批处理时间的计量器
    data_time = AverageMeter('Data', ':6.3f')  # 数据加载时间的计量器
    losses = AverageMeter('Loss', ':.4e')  # 损失的计量器
    acc = AverageMeter('Accuracy', ':7.3f')  # 准确率的计量器
    
    progress = ProgressMeter(
        len(train_loader),  # 总批次数
        [batch_time, data_time, losses, acc],  # 需要显示的计量器列表
        prefix=f"Epoch: [{epoch}]")  # 显示的前缀

    model.train()  # 切换到训练模式

    epoch_start = time.time()  # 记录一个epoch的开始时间
    end = epoch_start  # 初始化一个批次结束的时间点

    for i, (images, target) in enumerate(train_loader):  # 遍历数据加载器

        data_time.update(time.time() - end)  # 更新数据加载时间
    
        images = images.to(device,non_blocking=True)  # 将图像数据移动到指定设备
        target = target.to(device,non_blocking=True)  # 将目标数据移动到指定设备

        with amp.autocast():  # 自动混合精度上下文
            output = model(images)  # 计算模型输出
            loss = criterion(output, target)  # 计算损失
    
        losses.update(loss.item(), images.size(0))  # 更新损失计量器
        acc.update(accuracy(output, target), images.size(0))  # 更新准确率计量器

        optimizer.zero_grad()  # 清空梯度
        scaler.scale(loss).backward()  # 反向传播，计算梯度
        scaler.step(optimizer)  # 根据梯度更新模型参数
        scaler.update()  # 更新缩放器以进行下一轮
        batch_time.update(time.time() - end)  # 更新批处理时间
        end = time.time()  # 记录这一批次结束的时间点

        if i % 50 == 0 or i == len(train_loader)-1:  # 每50个批次或最后一个批次时显示进度
            progress.display(i)
    
    print(f"Epoch: [{epoch}] completed, elapsed time {time.time() - epoch_start:6.3f} seconds")  # 打印一个epoch的总耗时

    tensorboard.add_scalar('Loss/train', losses.avg, epoch)  # 将平均损失记录到tensorboard
    tensorboard.add_scalar('Accuracy/train', acc.avg, epoch)  # 将平均准确率记录到tensorboard
    return losses.avg, acc.avg  # 返回平均损失和平均准确率

def validate(val_loader, model, criterion, epoch):
    """
    Measure model performance across the val dataset
    """
    batch_time = AverageMeter('Time', ':6.3f')  # 验证过程中的批处理时间计量器
    losses = AverageMeter('Loss', ':.4e')  # 损失计量器
    acc = AverageMeter('Accuracy', ':7.3f')  # 准确率计量器
    
    progress = ProgressMeter(
        len(val_loader),  # 验证数据集的批次数
        [batch_time, losses, acc],  # 需要显示的计量器列表
        prefix='Val:   ')  # 显示的前缀，表示这是验证阶段

    model.eval()  # 切换模型到评估模式

    with torch.no_grad():  # 在此上下文中不计算梯度
        end = time.time()  # 记录开始时间
        for i, (images, target) in enumerate(val_loader):  # 遍历验证数据加载器
            images = images.to(device,non_blocking=True)  # 将图像数据移动到指定设备
            target = target.to(device,non_blocking=True)  # 将目标数据移动到指定设备
            # 计算模型输出
            with amp.autocast():  # 使用自动混合精度
                output = model(images)  # 获取模型对图像的输出
                loss = criterion(output, target)  # 计算损失
            # 更新损失和准确率计量器
            losses.update(loss.item(), images.size(0))
            acc.update(accuracy(output, target), images.size(0))
            # 更新批处理时间
            batch_time.update(time.time() - end)
            end = time.time()  # 记录当前时间为下一批次的开始时间
            if i % 10  == 0 or i == len(val_loader)-1:  # 每10个批次或最后一个批次时显示进度
                progress.display(i)

    tensorboard.add_scalar('Loss/val', losses.avg, epoch)  # 将平均损失记录到tensorboard
    tensorboard.add_scalar('Accuracy/val', acc.avg, epoch)  # 将平均准确率记录到tensorboard
    
    return losses.avg, acc.avg  # 返回平均损失和平均准确率
    
def save_checkpoint(state, is_best, filename='checkpoint.pth.tar', best_filename='model_best.pth.tar', labels_filename='labels.txt'):
    """
    Save a model checkpoint file, along with the best-performing model if applicable
    """
    model_dir = os.path.expanduser(model_path)  # 获取模型保存路径

    if not os.path.exists(model_dir):  # 如果模型保存路径不存在
        os.mkdir(model_dir)  # 创建该路径

    filename = os.path.join(model_dir, filename)  # 完整的检查点文件路径
    best_filename = os.path.join(model_dir, best_filename)  # 完整的最佳模型文件路径
    labels_filename = os.path.join(model_dir, labels_filename)  # 完整的标签文件路径
        
    torch.save(state, filename)  # 保存检查点
            
    if is_best:  # 如果是最佳模型
        shutil.copyfile(filename, best_filename)  # 将当前检查点复制为最佳模型文件
        print(f"saved best model to:  {best_filename}")
    else:
        print(f"saved checkpoint to:  {filename}")
        
    if state['epoch'] == 0:  # 如果是第一个epoch
        with open(labels_filename, 'w') as file:  # 打开标签文件进行写入
            for label in state['classes']:  # 遍历所有类别标签
                file.write(f"{label}\n")  # 写入每个标签
        print(f"saved class labels to:  {labels_filename}")  # 打印保存标签文件的信息
            
def main():
    global best_accuracy  # 使用全局变量来记录最佳准确率
    # 定义图像的标准化过程
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
    # 定义训练集的图像变换
    train_transforms = transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        normalize,
    ])

    # 定义验证集的图像变换
    val_transforms = transforms.Compose([
        transforms.Resize(224),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        normalize,
    ])
    # 加载训练集
    train_dataset = datasets.ImageFolder("./dataset/train", train_transforms)
    # 加载验证集
    val_dataset = datasets.ImageFolder("./dataset/val", val_transforms)
    # 创建训练集的数据加载器
    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=8, shuffle=True,
        num_workers=3, pin_memory=True)

    # 创建验证集的数据加载器
    val_loader = torch.utils.data.DataLoader(
        val_dataset, batch_size=16, shuffle=False,
        num_workers=3, pin_memory=True)
    # 初始化模型
    model = models.resnet18(pretrained=True)
    # 获取类别数
    num_classes = len(train_dataset.classes)
    # 替换模型的全连接层以匹配类别数
    model.fc = torch.nn.Linear(model.fc.in_features, num_classes)
    # 将模型移动到指定设备
    model = model.to(device)
    # 定义损失函数
    criterion = nn.CrossEntropyLoss()
    # 定义学习率、动量和权重衰减
    lr = 0.1
    momentum = 0.9 
    weight_decay = 1e-4
    # 初始化优化器
    optimizer = torch.optim.SGD(model.parameters(), lr,
                                momentum=momentum,
                                weight_decay=weight_decay)
    # 初始化梯度缩放器，用于自动混合精度训练
    scaler = amp.GradScaler()
    # 设置训练的总轮数
    epochs = 10
    for epoch in range(epochs):
        # 训练一个epoch
        train_loss, train_acc = train(train_loader, model, criterion, optimizer,scaler, epoch)
        # 在验证集上评估模型
        val_loss, val_acc = validate(val_loader, model, criterion, epoch)

        # 更新最佳准确率并保存模型
        is_best = val_acc > best_accuracy
        best_accuracy = max(val_acc, best_accuracy)

        # 打印训练和验证的损失及准确率
        print(f"=> Epoch {epoch}")
        print(f"  * Train Loss     {train_loss:.4e}")
        print(f"  * Train Accuracy {train_acc:.4f}")
        print(f"  * Val Loss       {val_loss:.4e}")
        print(f"  * Val Accuracy   {val_acc:.4f}{'*' if is_best else ''}")
        
        # 保存模型的检查点
        save_checkpoint({
            'epoch': epoch,
            'arch': "resnet18",
            'resolution': 224,
            'classes': train_dataset.classes,
            'num_classes': len(train_dataset.classes),
            'multi_label': False,
            'state_dict': model.state_dict(),
            'accuracy': {'train': train_acc, 'val': val_acc},
            'loss' : {'train': train_loss, 'val': val_loss},
            'optimizer' : optimizer.state_dict(),
        }, is_best)

if __name__ == '__main__':
    main()

离线推理脚本解析

脚本路径./omInfer/main.cpp

#include <cmath>  // 引入数学库
#include <dirent.h>  // 引入目录操作库
#include <string.h>  // 引入字符串操作库
#include <map>  // 引入map容器
#include "acllite_dvpp_lite/ImageProc.h"  // 引入图像处理头文件
#include "acllite_om_execute/ModelProc.h"  // 引入模型处理头文件

using namespace std;  // 使用标准命名空间
using namespace acllite;  // 使用acllite命名空间

int main()
{    
    vector<string> labels = { {"female"},{"male"}};  // 定义标签
    AclLiteResource aclResource;  // 创建AclLite资源对象
    bool ret = aclResource.Init();  // 初始化AclLite资源
    CHECK_RET(ret, LOG_PRINT("[ERROR] InitACLResource failed."); return 1);  // 检查初始化是否成功
    
    ImageProc imageProc;  // 创建图像处理对象
    ModelProc modelProc;  // 创建模型处理对象
    ret = modelProc.Load("../model/resnet18.om");  // 加载模型
    CHECK_RET(ret, LOG_PRINT("[ERROR] load model Resnet18.om failed."); return 1);  // 检查模型是否加载成功
    ImageData src = imageProc.Read("../data/8.jpg");  // 读取图像数据
    CHECK_RET(src.size, LOG_PRINT("[ERROR] ImRead image failed."); return 1);  // 检查图像是否读取成功
    
    ImageData dst;  // 定义处理后的图像数据
    ImageSize dsize(224, 224);  // 设置目标图像大小

    imageProc.Resize(src, dst, dsize);  // 调整图像大小    
    ret = modelProc.CreateInput(static_cast<void *>(dst.data.get()), dst.size);  // 创建模型输入
    CHECK_RET(ret, LOG_PRINT("[ERROR] Create model input failed."); return 1);  // 检查模型输入是否创建成功
    vector<InferenceOutput> inferOutputs;  // 定义推理输出
    ret = modelProc.Execute(inferOutputs);  // 执行模型推理
    CHECK_RET(ret, LOG_PRINT("[ERROR] model execute failed."); return 1);  // 检查模型推理是否执行成功

    uint32_t dataSize = inferOutputs[0].size;  // 获取推理输出数据大小
    // 从输出数据集中获取结果
    float* outData = static_cast<float*>(inferOutputs[0].data.get());  // 获取推理输出数据
    if (outData == nullptr) {
        LOG_PRINT("get result from output data set failed.");  // 如果获取失败，打印错误信息
        return 1;
    }
    int index = 0;  // 初始化最大值索引
    float max = 0;  // 初始化最大值
    for (uint32_t j = 0; j < dataSize / sizeof(float); ++j) {  // 遍历输出数据
        if (outData[j] > max){  // 如果当前值大于最大值
            max = outData[j];  // 更新最大值
            index = j;  // 更新最大值索引
        }
    }
    LOG_PRINT("[INFO] value[%lf] output[%s]", outData[index] , labels[index].c_str());  // 打印最大值和对应的标签
    outData = nullptr;  // 清空输出数据指针
    return 0;  // 程序正常结束
}

执行准备

• 登录开发板，进入代码目录，安装requirements

  (base) root@davinci-mini:~# cd EdgeAndRobotics/Samples/ClassficationRetrainingAndInfer/
  (base) root@davinci-mini:~/EdgeAndRobotics/Samples/ClassficationRetrainingAndInfer# pip install -r requirements.txt
  

  

• 安装PyTorch2.1.0、torchvision1.16.0
  

• 配置离线推理所需的环境变量

  (base) root@davinci-mini:~/EdgeAndRobotics/Samples/ClassficationRetrainingAndInfer# export DDK_PATH=/usr/local/Ascend/ascend-toolkit/latest
  (base) root@davinci-mini:~/EdgeAndRobotics/Samples/ClassficationRetrainingAndInfer# export NPU_HOST_LIB=$DDK_PATH/runtime/lib64/stub
  

  

• 安装acllite，参考链接

  Ascend/ACLLite - 码云 - 开源中国 (gitee.com)

模型训练

• 准备数据集

  (base) root@davinci-mini:~/EdgeAndRobotics/Samples/ClassficationRetrainingAndInfer# cd dataset/
  (base) root@davinci-mini:~/EdgeAndRobotics/Samples/ClassficationRetrainingAndInfer/dataset# wget https://obs-9be7.obs.cn-east-2.myhuaweicloud.com/wanzutao/gender.zip
  (base) root@davinci-mini:~/EdgeAndRobotics/Samples/ClassficationRetrainingAndInfer/dataset# unzip gender.zip
  

  

• 设置环境变量减小算子编译内存占用

  (base) root@davinci-mini:~/EdgeAndRobotics/Samples/ClassficationRetrainingAndInfer/dataset# export TE_PARALLEL_COMPILER=1
  (base) root@davinci-mini:~/EdgeAndRobotics/Samples/ClassficationRetrainingAndInfer/dataset# export MAX_COMPILE_CORE_NUMBER=1
  

• 运行训练脚本

  (base) root@davinci-mini:~/EdgeAndRobotics/Samples/ClassficationRetrainingAndInfer/dataset# cd ..
  (base) root@davinci-mini:~/EdgeAndRobotics/Samples/ClassficationRetrainingAndInfer# python3 main.py
  

  

离线推理

• 导出onnx模型

  python3 export.py
  

  

• 获取测试图片数据

  cd omInfer/data
  wget https://obs-9be7.obs.cn-east-2.myhuaweicloud.com/wanzutao/classfication/8.jpg
  

  
  

• 获取PyTorch框架的ResNet50模型（*.onnx），并转换为昇腾AI处理器能识别的模型（*.om）

  • 设置如下两个环境变量减少atc模型转换过程中使用的进程数，减小内存占用

  export TE_PARALLEL_COMPILER=1
  export MAX_COMPILE_CORE_NUMBER=1
  

  • 将导出的resnet18.onnx模型拷贝到model目录下

  cd ../model
  cp ../../resnet18.onnx ./
  

  • 获取AIPP配置文件

  wget https://obs-9be7.obs.cn-east-2.myhuaweicloud.com/wanzutao/classfication/aipp.cfg
  

  • 模型转换

  atc --model=resnet18.onnx --framework=5 --insert_op_conf=aipp.cfg --output=resnet18 --soc_version=Ascend310B4
  

  

• 编译样例源码

  cd ../scripts 
  bash sample_build.sh
  

  

• 执行以下脚本运行样例

  bash sample_run.sh