训练代码解析
脚本名称：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，参考链接
模型训练
-
准备数据集
(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
1065
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
