# eval example [0: 'voxels', 1: 'num_points', 2: 'coordinates', 3: 'rect'
# 4: 'Trv2c', 5: 'P2', 6: 'anchors', 7: 'anchors_mask'
# 8: 'image_idx', 9: 'image_shape']
0: 'voxels'
:体素信息,包括体素坐标、体素特征等。1: 'num_points'
:每个体素中包含的点云数量。2: 'coordinates'
:点云坐标信息,包括每个点的三维坐标值。3: 'rect'
:标注框的位置和大小信息,包括左上角和右下角的坐标值。4: 'Trv2c'
:相机和车辆坐标系之间的变换矩阵。5: 'P2'
:相机投影矩阵。6: 'anchors'
:锚框信息,包括锚框的位置、大小等。7: 'anchors_mask'
:锚框掩码,用于指示哪些锚框需要进行目标检测。8: 'image_idx'
:图像索引,用于指示当前数据属于哪张图像。9: 'image_shape'
:图像尺寸,包括图像的宽度和高度。
首先看的是训练函数,下面是训练函数传入了一些基本参数 :
config_path
:模型配置文件的路径;model_dir
:模型存储的目录,训练过程中的模型检查点和最终的模型将存储在此目录下;result_path
:训练结果的存储路径,默认为None
,表示不存储训练结果;create_folder
:是否创建model_dir
和result_path
指定的目录,默认为False
;display_step
:每训练多少步打印一次训练信息,默认为50
;summary_step
:每训练多少步记录一次训练摘要信息,默认为5
;pickle_result
:训练结果是否序列化存储,默认为True
。 函数并没有返回值,而是将训练过程中的模型检查点和最终的模型保存在model_dir
目录下,如果指定了result_path
,还会将训练结果存储到该路径
def train(config_path,
model_dir,
result_path=None,
create_folder=False,
display_step=50,
summary_step=5,
pickle_result=True):
if create_folder:
if pathlib.Path(model_dir).exists():
model_dir = torchplus.train.create_folder(model_dir)
# parents = True表示如果目录的上级目录不存在,则将同时创建上级目录;exist_ok = True表示如果目录已经存在,则不会抛出异常。
model_dir = pathlib.Path(model_dir)
model_dir.mkdir(parents=True, exist_ok=True)
eval_checkpoint_dir = model_dir / 'eval_checkpoints'
eval_checkpoint_dir.mkdir(parents=True, exist_ok=True)
if result_path is None:
result_path = model_dir / 'results'
config_file_bkp = "pipeline.config" # 备份原始的配置文件
config = pipeline_pb2.TrainEvalPipelineConfig() # 存储模型的训练和评估配置
with open(config_path, "r") as f:
proto_str = f.read()
text_format.Merge(proto_str, config) # 将proto_str中的文本格式的内容解析为config实例对应的protobuf消息类型
shutil.copyfile(config_path, str(model_dir / config_file_bkp)) # 原始配置文件复制到指定的备份路径下
input_cfg = config.train_input_reader # 训练数据输入配置
eval_input_cfg = config.eval_input_reader # 评估数据输入配置
model_cfg = config.model.second # 模型配置
train_cfg = config.train_config # 训练配置
注解:
1
.pipeline_pb2
是protobuf生成的Python模块,其中定义了一系列protobuf消息类型,包括用于存储模型训练和评估配置的TrainEvalPipelineConfig
消息类型。2.
text_format
模块,它提供了将protobuf消息类型与文本格式之间进行相互转换的函数。通过将配置文件解析为config
实例,可以方便地对其进行修改和保存。
class_names = list(input_cfg.class_names)
#########################
# Build Voxel Generator
#########################
voxel_generator = voxel_builder.build(model_cfg.voxel_generator)
#########################
# Build Target Assigner
#########################
bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]]
box_coder = box_coder_builder.build(model_cfg.box_coder)
target_assigner_cfg = model_cfg.target_assigner
target_assigner = target_assigner_builder.build(target_assigner_cfg, bv_range, box_coder)
首先,从训练数据输入配置中获取类别名称,转换为列表类型的
class_names
变量。然后,使用
voxel_builder
模块中的build
函数根据模型配置中的体素生成器配置,构建相应的体素生成器voxel_generator
。接着,从模型配置中获取点云范围并计算出对应的BEV范围
bv_range
,并使用box_coder_builder
模块中的build
函数构建相应的box_coder
。最后,从模型配置中获取目标分配器配置,并使用
target_assigner_builder
模块中的build
函数构建相应的target_assigner
,以便后续在训练和评估中使用。
下面代码实现了模型构建的一些准备工作,包括根据输入配置中的类别名称获取类别列表,构建体素生成器和目标分配器以及设置优化器和损失函数。
######################
# Build NetWork
######################
center_limit_range = model_cfg.post_center_limit_range
# net = second_builder.build(model_cfg, voxel_generator, target_assigner)
net = second_builder.build(model_cfg, voxel_generator, target_assigner, input_cfg.batch_size)
net.cuda()
# net_train = torch.nn.DataParallel(net).cuda()
print("num_trainable parameters:", len(list(net.parameters())))
# for n, p in net.named_parameters():
# print(n, p.shape)
######################
# Build Optimizer
######################
# we need global_step to create lr_scheduler, so restore net first.
# 我们需要globalstep来创建lrscheduler,所以先恢复net。
torchplus.train.try_restore_latest_checkpoints(model_dir, [net]) # 模型恢复
gstep = net.get_global_step() - 1
optimizer_cfg = train_cfg.optimizer
if train_cfg.enable_mixed_precision:
net.half()
net.metrics_to_float()
net.convert_norm_to_float(net)
optimizer = optimizer_builder.build(optimizer_cfg, net.parameters())
if train_cfg.enable_mixed_precision:
loss_scale = train_cfg.loss_scale_factor
mixed_optimizer = torchplus.train.MixedPrecisionWrapper(optimizer, loss_scale)
else:
mixed_optimizer = optimizer
# must restore optimizer AFTER using MixedPrecisionWrapper
# 必须在使用MixedPrecisionWrapper之后恢复优化器
torchplus.train.try_restore_latest_checkpoints(model_dir, [mixed_optimizer])
lr_scheduler = lr_scheduler_builder.build(optimizer_cfg, optimizer, gstep) # 优化器
# 根据训练配置中的enable_mixed_precision参数来设置浮点数类型。
if train_cfg.enable_mixed_precision:
float_dtype = torch.float16
else:
float_dtype = torch.float32
#损失函数是multi-loss的形式,包含了回归损失和分类损失,分类损失采用的是focal-loss
模型的保存与恢复(try_restore_latest_checkpoints):【深度学习笔记(十六)】之tensorflow2中模型的保存与恢复_tensorflow2保存模型_开发小鸽的博客-CSDN博客
这段代码根据训练配置中的
enable_mixed_precision
参数来构建优化器,并在需要时将模型转换为混合精度模式。如果
enable_mixed_precision
为True
,则将模型转换为半精度模式(net.half()
),并将模型中的度量指标(metrics)转换为浮点数(net.metrics_to_float()
)。如果模型中存在归一化层,则还需要将其转换为浮点数(net.convert_norm_to_float(net)
)。然后,使用优化器配置和网络参数构建优化器(optimizer_builder.build(optimizer_cfg, net.parameters())
)。如果
enable_mixed_precision
为True
,则使用MixedPrecisionWrapper
将优化器转换为混合精度优化器,并指定损失缩放因子(train_cfg.loss_scale_factor
)。如果enable_mixed_precision
为False
,则直接使用构建好的优化器。最终返回构建好的混合精度优化器或普通优化器。该代码段将用于在训练模型时构建优化器
######################
# Prepare Input 预处理
######################
dataset = input_reader_builder.build(
input_cfg,
model_cfg,
training=True,
voxel_generator=voxel_generator,
target_assigner=target_assigner)
eval_dataset = input_reader_builder.build(
eval_input_cfg,
model_cfg,
training=False,
voxel_generator=voxel_generator,
target_assigner=target_assigner)
# 设置PyTorch多进程数据加载器的初始化函数,用于设置每个工作进程的随机数种子。
def _worker_init_fn(worker_id):
time_seed = np.array(time.time(), dtype=np.int32)
np.random.seed(time_seed + worker_id)
print(f"WORKER {worker_id} seed:", np.random.get_state()[1][0])
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=input_cfg.batch_size,
shuffle=True,
num_workers=input_cfg.num_workers,
pin_memory=False,
collate_fn=merge_second_batch,
worker_init_fn=_worker_init_fn)
eval_dataloader = torch.utils.data.DataLoader(
eval_dataset,
batch_size=eval_input_cfg.batch_size,
shuffle=False,
num_workers=eval_input_cfg.num_workers,
pin_memory=False,
collate_fn=merge_second_batch)
data_iter = iter(dataloader)
使用
input_reader_builder
对象的build
函数构建了训练数据集dataset
和测试数据集eval_dataset
,具体参数如下:
input_cfg
和eval_input_cfg
:数据读取相关的配置信息,包括数据集路径、文件格式、采样器等参数。model_cfg
:模型配置信息,包括模型架构、输入输出通道数、损失函数等参数。training
:表示当前是训练阶段还是测试阶段。如果为True,则表示训练阶段,数据集将包含标注信息,用于计算损失函数;否则表示测试阶段,数据集将不包含标注信息,用于模型的预测和评估。voxel_generator
:体素化器对象,用于将点云数据转换为体素化表示。target_assigner
:目标分配器对象,用于将标注信息与体素化后的数据进行匹配,生成用于计算损失函数的训练目标。 数据集的构建过程通常包括以下步骤:
- 使用输入读取器(input_reader)从硬盘或其他存储介质中读取数据。
- 使用体素化器(voxel_generator)将点云数据转换为体素化表示。
- 使用目标分配器(target_assigner)将标注信息与体素化后的数据进行匹配,生成用于计算损失函数的训练目标。
- 将训练数据集(dataset)和测试数据集(eval_dataset)返回给调用者。
定义了两个PyTorch数据加载器,即
dataloader
和eval_dataloader
,并使用iter
函数将dataloader
转换为数据迭代器data_iter。
具体参数如下:
dataset
和eval_dataset
:表示训练数据集和测试数据集。这两个参数通常是继承自torch.utils.data.Dataset
类的自定义数据集,用于读取和预处理输入数据。batch_size
:表示批次大小。shuffle
:表示是否打乱数据集。num_workers
:表示使用的工作进程数。通过使用多进程加载数据可以加速数据加载和预处理过程,提高模型训练速度。但是,使用过多的进程可能会导致系统资源的瓶颈,从而影响训练效率。pin_memory
:表示是否将数据加载到GPU内存中。如果设置为True,则可以加速数据传输,提高训练效率。但是,如果GPU内存不足,可能会导致训练失败。collate_fn
:表示如何合并不同的样本数据。在处理不同大小的样本时,需要将它们合并为一个批次,并进行填充等操作。collate_fn
函数通常是一个自定义函数,用于实现不同样本的合并方法。worker_init_fn
:表示在每个工作进程启动时调用的函数。该函数通常用于设置每个工作进程的随机数种子,以保证训练过程的随机性和可复现性。
######################
# Training
######################
log_path = model_dir / 'log.txt'
logf = open(log_path, 'a')
logf.write(proto_str)
logf.write("\n")
summary_dir = model_dir / 'summary'
summary_dir.mkdir(parents=True, exist_ok=True)
writer = SummaryWriter(str(summary_dir)) # 记录模型训练过程中的损失函数、精度、梯度等信息
total_step_elapsed = 0
remain_steps = train_cfg.steps - net.get_global_step()
t = time.time()
ckpt_start_time = t # 表示每个检查点开始的时间
total_loop = train_cfg.steps // train_cfg.steps_per_eval + 1
# total_loop = remain_steps // train_cfg.steps_per_eval + 1
clear_metrics_every_epoch = train_cfg.clear_metrics_every_epoch
if train_cfg.steps % train_cfg.steps_per_eval == 0:
total_loop -= 1
mixed_optimizer.zero_grad() #清空评估指标。
for _ in range(total_loop):
# 判断当前循环是否为最后一个循环
if total_step_elapsed + train_cfg.steps_per_eval > train_cfg.steps:
steps = train_cfg.steps % train_cfg.steps_per_eval
else:
steps = train_cfg.steps_per_eval
for step in range(steps):
lr_scheduler.step() # 更新学习率
try:
example = next(data_iter) # 获取下一个样本数据
except StopIteration: # 代表迭代完成,抛出异常StopIteration
print("end epoch")
if clear_metrics_every_epoch: # 清空评估指标
net.clear_metrics()
data_iter = iter(dataloader) # 重置数据迭代器
example = next(data_iter)
example_torch = example_convert_to_torch(example, float_dtype)
batch_size = example["anchors"].shape[0]
example_tuple = list(example_torch.values())
example_tuple[11] = torch.from_numpy(example_tuple[11])
example_tuple[12] = torch.from_numpy(example_tuple[12])
assert 13 == len(example_tuple), "something write with training input size!"
- 从获取到的数据样本
example
中获取锚框的数量,并将其赋值给batch_size
变量。- 将获取到的数据样本
example_torch
中的numpy数组转换为PyTorch的张量格式。其中,example_tuple[11]和example_tuple[12]分别对应于masks和classes,需要使用torch.from_numpy
方法转换为PyTorch的张量。- 检查转换后的数据样本
example_tuple
的长度是否为13,如果不是则说明输入大小有误,抛出异常提示信息。
pillar_x = example_tuple[0][:, :, 0].unsqueeze(0).unsqueeze(0)
pillar_y = example_tuple[0][:, :, 1].unsqueeze(0).unsqueeze(0)
pillar_z = example_tuple[0][:, :, 2].unsqueeze(0).unsqueeze(0)
pillar_i = example_tuple[0][:, :, 3].unsqueeze(0).unsqueeze(0)
num_points_per_pillar = example_tuple[1].float().unsqueeze(0)
上面代码用于从数据样本中获取点云数据,将其转换为柱状体(pillar)表示法,并将其转换为PyTorch的张量格式,具体实现了以下功能:
- 从数据样本中获取点云数据,具体包括柱状体的坐标和强度信息。其中,example_tuple[0]对应于坐标和强度信息,example_tuple[1]对应于每个柱状体中点的数量。
- 提取柱状体中的x、y、z坐标和强度信息,并通过unsqueeze方法扩展维度,使得它们的shape为[1, 1, N, V],其中N为柱状体数量,V为每个柱状体中点的数量。
- 将每个柱状体中点的数量转换为float类型,并通过unsqueeze方法扩展维度,使得其shape为[1, 1, N]。
# 提取x、y坐标信息
coors_x = example_tuple[2][:, 3].float()
coors_y = example_tuple[2][:, 2].float()
x_sub = coors_x.unsqueeze(1) * 0.16 + 0.08
y_sub = coors_y.unsqueeze(1) * 0.16 - 39.6
ones = torch.ones([1, 100], dtype=torch.float32, device=pillar_x.device)
x_sub_shaped = torch.mm(x_sub, ones).unsqueeze(0).unsqueeze(0)
y_sub_shaped = torch.mm(y_sub, ones).unsqueeze(0).unsqueeze(0)
num_points_for_a_pillar = pillar_x.size()[3]
mask = get_paddings_indicator(num_points_per_pillar, num_points_for_a_pillar, axis=0)
mask = mask.permute(0, 2, 1)
mask = mask.unsqueeze(1)
mask = mask.type_as(pillar_x)
这段代码用于构造一个掩码张量,用于标记每个柱状体中哪些点是有效的,具体实现了以下功能:
- 获取每个柱状体中点的数量,以及每个柱状体中最大点的数量,分别赋值给
num_points_per_pillar
和num_points_for_a_pillar
变量。- 调用
get_paddings_indicator
函数生成掩码张量,其中axis=0
表示在第0个维度上进行padding,掩码张量的shape为[num_points_per_pillar, num_points_for_a_pillar]。- 将掩码张量沿着第0个维度进行转置,使得它的shape变为[N, num_points_for_a_pillar, num_points_per_pillar],其中N为柱状体的数量。
- 在掩码张量的第1个维度上添加一个维度,使得它的shape变为[N, 1, num_points_for_a_pillar, num_points_per_pillar]。
- 将掩码张量的数据类型转换为与pillar_x相同的数据类型。
coors = example_tuple[2]
anchors = example_tuple[6]
labels = example_tuple[8]
reg_targets = example_tuple[9]
input = [pillar_x, pillar_y, pillar_z, pillar_i, num_points_per_pillar,
x_sub_shaped, y_sub_shaped, mask, coors, anchors, labels, reg_targets]
ret_dict = net(input)
assert 10 == len(ret_dict), "something write with training output size!"
coors |
表示每个anchor在原始图片上的坐标,一般用于计算anchor与ground truth box之间的IoU(交并比)。 |
anchors |
表示生成的anchor的坐标,是目标检测中用于表示可能存在物体的矩形框。 |
labels |
表示每个anchor所对应的物体类别,是目标检测中的分类标签。 |
reg_targets | 表示每个anchor与其对应的ground truth box之间的偏移量,是目标检测中的回归目标 |
ret_dict是一个字典类型,包含了目标检测模型的输出信息,其中有10个键值对表示10个不同的输出,包括:
cls_preds | 预测的物体类别概率矩阵; |
box_preds | 预测的物体框坐标矩阵; |
dir_cls_preds | 预测的物体方向角度概率矩阵; |
pillar_features | :pillar特征矩阵 |
seg_preds | 预测的点云分割结果矩阵; |
points | 网络输入的点云数据; |
coors | anchor坐标矩阵; |
anchors | anchor框坐标矩阵; |
labels | anchor框对应的物体类别标签矩阵; |
reg_targets | anchor框与ground truth box之间的回归目标矩阵。 |
# 从模型输出的结果ret_dict中获取分类损失、定位损失、正样本分类损失、负样本分类损失、正样本数、负样本数、分类预测结果等信息。
cls_preds = ret_dict[5]
loss = ret_dict[0].mean()
cls_loss_reduced = ret_dict[7].mean()
loc_loss_reduced = ret_dict[8].mean()
cls_pos_loss = ret_dict[3]
cls_neg_loss = ret_dict[4]
loc_loss = ret_dict[2]
cls_loss = ret_dict[1]
dir_loss_reduced = ret_dict[6]
cared = ret_dict[9]
labels = example_tuple[8]
# 如果开启了混合精度训练,则对损失进行缩放。
if train_cfg.enable_mixed_precision:
loss *= loss_scale
loss.backward() # 反向传播
torch.nn.utils.clip_grad_norm_(net.parameters(), 10.0) # 对梯度进行裁剪 第一个参数是模型的参数,第二个参数是裁剪的范数,即梯度值的最大范数,
mixed_optimizer.step() # 更新模型参数
mixed_optimizer.zero_grad() # 清空模型梯度的函数
net.update_global_step() # 更新全局步骤计数器
net_metrics = net.update_metrics(cls_loss_reduced,loc_loss_reduced,
cls_preds,labels,cared)
# 在更新训练指标时,会调用目标检测模型的update_metrics()方法,该方法会计算网络的预测结果与真实标签之间的精度、召回率、F1值等指标
cls_preds |
表示预测的物体类别概率矩阵 |
loss |
表示目标检测模型的总损失 |
cls_loss_reduced |
表示分类损失的均值 |
loc_loss_reduced |
表示定位损失的均值 |
cls_pos_loss |
表示正样本分类损失 |
cls_neg_loss |
表示负样本分类损失 |
loc_loss |
表示定位损失 |
cls_loss |
表示分类损失 |
dir_loss_reduced |
表示方向角度损失的均值 |
cared |
表示可用于计算损失的anchor的掩码矩阵 |
labels |
表示anchor对应的物体类别标签矩阵 |
net_parameters |
() 表示网络参数 |
mixed_optimizer
是一个混合精度优化器,通常是torch.cuda.amp.GradScaler()
类的实例。混合精度优化器的原理是将模型的参数和梯度分别转换为低精度半精度(FP16)和高精度(FP32)的数据类型进行计算,可以加速模型的训练,同时减少GPU显存的使用量。在更新模型参数之前,需要调用zero_grad()
函数清空之前的梯度信息。在使用混合精度训练时,需要对损失值进行缩放,即将损失值乘以一个缩放因子,通常是loss_scale = mixed_optimizer.get_scale()
,然后再进行反向传播计算梯度。更新模型参数时,需要调用mixed_optimizer.step()
函数,该函数会根据梯度计算出参数的更新量,并更新模型参数。
step_time = (time.time() - t)
t = time.time()
metrics = {}
num_pos = int((labels > 0)[0].float().sum().cpu().numpy()) # 正样本
num_neg = int((labels == 0)[0].float().sum().cpu().numpy()) # 负样本
num_anchors = int(example_tuple[7][0].sum())
global_step = net.get_global_step()
# 训练指标显示
if global_step % display_step == 0:
loc_loss_elem = [
float(loc_loss[:, :, i].sum().detach().cpu().numpy() /
batch_size) for i in range(loc_loss.shape[-1]) # 每个定位损失元素的平均值
]
metrics["step"] = global_step
metrics["steptime"] = step_time
metrics.update(net_metrics)
metrics["loss"] = {}
metrics["loss"]["loc_elem"] = loc_loss_elem
metrics["loss"]["cls_pos_rt"] = float(cls_pos_loss.detach().cpu().numpy())
metrics["loss"]["cls_neg_rt"] = float(cls_neg_loss.detach().cpu().numpy())
第一段记录训练时间。
第二段计算当前batch中anchor数量的变量
第三段训练指标显示和更新在训练指标的更新过程中,
首先计算了每个定位损失元素的平均值,这里使用了列表推导式和
detach()
函数将定位损失矩阵中的元素值提取出来并转换为numpy数组,然后除以batch_size得到平均值。接着,将训练步骤数、训练时间、网络指标、损失等信息保存在
metrics
字典中,以便后续对模型的训练效果进行监控和分析。其中,update()
函数用于将net_metrics
中的指标更新到metrics
字典中。最后,将定位损失、正样本分类损失和负样本分类损失的平均值保存在
metric
s["loss"]
字典中,用于后续的可视化和分析。
# 如果目标检测模型中包含方向分类器,则将方向损失的值也加入到metrics["loss"]字典中。
if model_cfg.use_direction_classifier:
metrics["loss"]["dir_rt"] = float(dir_loss_reduced.detach().cpu().numpy())
metrics["num_vox"] = int(example_tuple[0].shape[0])
metrics["num_pos"] = int(num_pos)
metrics["num_neg"] = int(num_neg)
metrics["num_anchors"] = int(num_anchors)
metrics["lr"] = float(mixed_optimizer.param_groups[0]['lr'])
metrics["image_idx"] = example_tuple[11][0]
flatted_metrics = flat_nested_json_dict(metrics)
flatted_summarys = flat_nested_json_dict(metrics, "/")
# flat_nested_json_dict()函数可以将嵌套的字典结构进行扁平化处理,将所有键值对都保存在一个字典中,方便后续对指标进行可视化和分析。
for k, v in flatted_summarys.items():
if isinstance(v, (list, tuple)):
v = {str(i): e for i, e in enumerate(v)}
writer.add_scalars(k, v, global_step)
else:
writer.add_scalar(k, v, global_step)
metrics_str_list = []
for k, v in flatted_metrics.items():
if isinstance(v, float):
metrics_str_list.append(f"{k}={v:.3}")
elif isinstance(v, (list, tuple)):
if v and isinstance(v[0], float):
v_str = ', '.join([f"{e:.3}" for e in v])
metrics_str_list.append(f"{k}=[{v_str}]")
else:
metrics_str_list.append(f"{k}={v}")
else:
metrics_str_list.append(f"{k}={v}")
这段代码中,首先遍历
flatted_summarys
中的键值对,将结果保存在k
和v
中。
- 如果
v
是列表或元组类型,则将其转换为字典,键为字符串类型的索引,值为列表中的元素值。然后,将k
和v
以及当前的训练步骤数作为参数,调用add_scalars()
函数将指标写入TensorBoard中。- 如果
v
是其他类型,则直接调用add_scalar()
函数将指标写入TensorBoard中。接着,遍历
flatted_metrics
中的键值对,将结果保存在k
和v
中。如果
v
是浮点型,则将其保留3位小数,并将k
和v
拼接成字符串,保存在metrics_str_list
中。如果
v
是列表或元组类型,则判断列表中的元素是否为浮点型,如果是,则将每个元素保留3位小数并拼接成字符串,保存在
metrics_str_list
中。如果列表中的元素不是浮点型,则直接将
k
和v
拼接成字符串,并保存在metrics_str_list
中。最后,将
metrics_str_list
中的所有字符串元素使用逗号拼接成一个字符串,作为当前batch的指标信息,并打印输出。
ckpt_elasped_time = time.time() - ckpt_start_time # 计算当前检查点(checkpoint)的训练时间
if ckpt_elasped_time > train_cfg.save_checkpoints_secs: # 为真则将当前模型的状态保存到磁盘上
torchplus.train.save_models(model_dir, [net, optimizer], net.get_global_step())
ckpt_start_time = time.time()
如果超时了,则将当前模型的状态保存到磁盘上。具体地,调用
torchplus.train.save_models()
函数,将当前模型的网络结构、优化器状态和全局步数等信息保存到指定的模型目录model_dir
下。最后
,将当前时间保存在
ckpt_start_time
中,作为下一个检查点的开始时间。
except Exception as e:
torchplus.train.save_models(model_dir, [net, optimizer], net.get_global_step())
logf.close()
raise e
与上面的try对应,如果程序出现异常,则调用
torchplus.train.save_models()
函数将当前模型的状态保存到指定的模型目录model_dir
下,以便后续恢复模型状态,避免训练过程中的损失。然后关闭日志,抛出异常。
def get_paddings_indicator(actual_num, max_num, axis=0):
这个函数的作用是生成一个指示张量,用于指示每个样本在填充后的长度。具体来说,它接受3个参数:
actual_num
: 一个长度为batch_size
的一维张量,表示每个样本在填充前的实际长度。max_num
: 所有样本中最大的长度。axis
: 指示应该在哪个维度上添加填充指示符。默认值为0,表示在最外层添加。
actual_num = torch.unsqueeze(actual_num, axis + 1) # 增加一个维度,从一维变成二维
max_num_shape = [1] * len(actual_num.shape) # 创建一个形状为[1,1,...,1]的张量
max_num_shape[axis + 1] = -1 # 在axis+1维度上为 -1
max_num = torch.arange(max_num, dtype=torch.int, device=actual_num.device).view(max_num_shape)
# tiled_actual_num : [N, M, 1]
# tiled_actual_num : [[3,3,3,3,3], [4,4,4,4,4], [2,2,2,2,2]]
# title_max_num : [[0,1,2,3,4], [0,1,2,3,4], [0,1,2,3,4]]
paddings_indicator = actual_num.int() > max_num # 其中每个元素的值为0或1,表示对应位置是否为填充位置。
# paddings_indicator shape : [batch_size, max_num]
return paddings_indicator
示例:假设有一个形状为
(3,)
的张量actual_num
,其中包含了三个样本的实际序列长度,分别为3、4、2。假设axis=0
,max_num=5
,则根据上述代码,可以得到以下结果:actual_num = torch.tensor([3, 4, 2]) actual_num = torch.unsqueeze(actual_num, axis+1) # actual_num.shape: (3, 1) max_num_shape = [1] * len(actual_num.shape) max_num_shape[axis+1] = -1 # max_num_shape: [1, -1] max_num = torch.arange(max_num, dtype=torch.int, device=actual_num.device).view(max_num_shape) # max_num: tensor([[0, 1, 2, 3, 4], # [0, 1, 2, 3, 4], # [0, 1, 2, 3, 4]]) paddings_indicator = actual_num.int() > max_num # paddings_indicator: tensor([[False, False, False, True, True], # [False, False, False, False, True], # [False, False, True, True, True]])
其中,
actual_num
张量的形状从(3,)
变为(3,1)
,max_num
张量的形状为(3,5)
,paddings_indicator
张量的形状为(3,5)
。可以看到,对于每个样本,填充位置的值为True
,有效位置的值为False
。这样,可以方便地将不同长度的序列按照最大序列长度进行批量化处理。
def _get_pos_neg_loss(cls_loss, labels):
该函数用于计算二分类问题中的正负样本损失。输入参数为
cls_loss
(二分类损失函数的输出)和labels
(真实标签),输出为正样本和负样本的损失。
def _get_pos_neg_loss(cls_loss, labels):
# cls_loss: [N, num_anchors, num_class]
# labels: [N, num_anchors]
batch_size = cls_loss.shape[0]
# 判断是否是一维或者二维张量
if cls_loss.shape[-1] == 1 or len(cls_loss.shape) == 2:
cls_pos_loss = (labels > 0).type_as(cls_loss) * cls_loss.view(
batch_size, -1)
cls_neg_loss = (labels == 0).type_as(cls_loss) * cls_loss.view(
batch_size, -1)
cls_pos_loss = cls_pos_loss.sum() / batch_size # 平均损失
cls_neg_loss = cls_neg_loss.sum() / batch_size
else:
cls_pos_loss = cls_loss[..., 1:].sum() / batch_size
cls_neg_loss = cls_loss[..., 0].sum() / batch_size
return cls_pos_loss, cls_neg_loss
输入参数为
cls_loss
(二分类损失函数的输出)和labels
(真实标签),输出为正样本和负样本的损失。 具体实现如下:首先,获取
cls_loss
张量的形状中的批量大小,并判断cls_loss
张量是否为一维或二维张量。如果是一维或二维张量,则说明每个样本只有一个类别预测值,此时需要将cls_loss
张量变成二维张量。然后,根据labels
张量中元素的值是否为1或0,将cls_loss
张量中对应位置上的元素分为正样本和负样本。使用type_as()
函数将labels
张量的类型转换成cls_loss
张量的类型,并将正样本和负样本的损失分别除以批量大小后返回,即得到正样本和负样本的平均损失。如果
cls_loss
张量不是一维或二维张量,则说明每个样本有多个类别预测值。此时,第0列代表负样本,其他列代表正样本。因此,可以通过[..., 1:]
获取正样本的损失,通过[..., 0]
获取负样本的损失。然后,将正样本和负样本的损失分别除以批量大小后返回,即得到正样本和负样本的平均损失。其中,
type_as()
函数用于将一个张量的类型转换为另一个张量的类型。view()
函数用于改变张量的形状,其中参数-1
表示根据张量的总大小和其他维度的大小自动推算该维度的大小。
def _flat_nested_json_dict(json_dict, flatted, sep=".", start=""):
for k, v in json_dict.items():
if isinstance(v, dict):
_flat_nested_json_dict(v, flatted, sep, start + sep + k)
else:
flatted[start + sep + k] = v
该函数用于将嵌套的JSON字典展平成一维字典。
输入参数为
json_dict
(嵌套的JSON字典)、flatted
(展平后的一维字典)和sep
(键名分隔符,默认为".")、start
(键名前缀,默认为空字符串)。输出为展平后的一维字典。具体实现如下: 首先,遍历
json_dict
字典的所有键值对如果当前值是字典,则递归调用
_flat_nested_json_dict()
函数将该子字典展平。否则,将当前键名和键值拼接成键值对,并以键名前缀为前缀加入到
flatted
字典中# 嵌套字典 nested_dict = { "a": {"b": 1, "c": 2}, "d": 3 } # 将嵌套字典展平为一维字典 flattened_dict = _flat_nested_json_dict(nested_dict, {}) print(flattened_dict) ###{'a.b': 1, 'a.c': 2, 'd': 3}
def flat_nested_json_dict(json_dict, sep=".") -> dict:
"""flat a nested json-like dict. this function make shadow copy.
"""
flatted = {}
for k, v in json_dict.items():
if isinstance(v, dict):
_flat_nested_json_dict(v, flatted, sep, k)
else:
flatted[k] = v
return flatted
def example_convert_to_torch(example, dtype=torch.float32, device=None) -> dict:
device = device or torch.device("cuda:0")
example_torch = {}
float_names = ["voxels", "anchors", "reg_targets", "reg_weights", "bev_map", "rect", "Trv2c", "P2"]
for k, v in example.items():
if k in float_names:
example_torch[k] = torch.as_tensor(v, dtype=dtype, device=device)
elif k in ["coordinates", "labels", "num_points"]:
example_torch[k] = torch.as_tensor(v, dtype=torch.int32, device=device)
elif k in ["anchors_mask"]:
example_torch[k] = torch.as_tensor(v, dtype=torch.uint8, device=device)
# torch.uint8 is now deprecated, please use a dtype torch.bool instead
else:
example_torch[k] = v
return example_torch
将输入数据转换为PyTorch张量的函数。输入是一个数据字典,其中每个键值对表示一个特征和对应的数值,输出是一个PyTorch张量字典,其中每个键值对表示一个特征和对应的张量。 具体来说,该函数首先创建一个空的PyTorch张量字典。然后遍历输入的数据字典,对于每个键值对,将其值转换为对应的PyTorch张量,并将该张量加入到输出的张量字典中。 该函数的参数说明如下:
example
:一个数据字典,其中每个键值对表示一个特征和对应的数值。dtype
:要转换成的PyTorch张量的数据类型,默认为torch.float32
。device
:要将PyTorch张量放置在哪个设备上,默认为None
,表示使用当前默认设备。对于不同的键名会转换成不同类型的张量
def _predict_kitti_to_file(net,
example,
result_save_path,
class_names,
center_limit_range=None,
lidar_input=False):
用于预测KITTI数据集上物体检测结果并将结果保存到文件的函数。该函数的输入包括:
net
:一个PyTorch模型,用于进行物体检测。example
:一个数据字典,包含了待检测的点云、标签等信息。result_save_path
:保存检测结果的文件路径。class_names
:一个列表,表示待检测的物体类别名称。center_limit_range
:一个长度为3的列表或元组,表示点云数据的范围限制。lidar_input
:一个布尔值,表示是否使用激光雷达数据作为模型的输入。
对检测结果进行处理
batch_image_shape = example['image_shape']
batch_imgidx = example['image_idx']
predictions_dicts = net(example) # 预测
for i, preds_dict in enumerate(predictions_dicts):
image_shape = batch_image_shape[i] # 图像尺寸
img_idx = preds_dict["image_idx"] # 图像索引
# 判断当前检测结果对应的2D边界框是否为None
if preds_dict["bbox"] is not None:
# 2D边界框、3D边界框、置信度和类别等信息
box_2d_preds = preds_dict["bbox"].data.cpu().numpy()
box_preds = preds_dict["box3d_camera"].data.cpu().numpy()
scores = preds_dict["scores"].data.cpu().numpy()
box_preds_lidar = preds_dict["box3d_lidar"].data.cpu().numpy()
# write pred to file 3D边界框的坐标顺序调整为hwl格式(即高度、宽度、长度),并将其赋值给box_preds变量。
box_preds = box_preds[:, [0, 1, 2, 4, 5, 3, 6]] # lhw->hwl(label file format)
label_preds = preds_dict["label_preds"].data.cpu().numpy() # 预测结果
# label_preds = np.zeros([box_2d_preds.shape[0]], dtype=np.int32)
result_lines = [] # 结果
for box, box_lidar, bbox, score, label in zip(
box_preds, box_preds_lidar, box_2d_preds, scores,
label_preds):
if not lidar_input:
if bbox[0] > image_shape[1] or bbox[1] > image_shape[0]:
continue
if bbox[2] < 0 or bbox[3] < 0:
continue
# print(img_shape)
if center_limit_range is not None:
limit_range = np.array(center_limit_range)
if (np.any(box_lidar[:3] < limit_range[:3])
or np.any(box_lidar[:3] > limit_range[3:])):
continue
bbox[2:] = np.minimum(bbox[2:], image_shape[::-1])
bbox[:2] = np.maximum(bbox[:2], [0, 0])
result_dict = {
'name': class_names[int(label)],
'alpha': -np.arctan2(-box_lidar[1], box_lidar[0]) + box[6],
'bbox': bbox,
'location': box[:3],
'dimensions': box[3:6],
'rotation_y': box[6],
'score': score,
}
result_line = kitti.kitti_result_line(result_dict)
result_lines.append(result_line)
- 使用
zip()
函数同时遍历box_preds
、box_preds_lidar
、box_2d_preds
、scores
和label_preds
数组,依次获取每个检测结果对应的3D边界框、激光雷达坐标系下的3D边界框、2D边界框、置信度和类别信息。- 如果当前输入数据是图像,获取当前检测结果的2D边界框在图像坐标系下的左上角和右下角坐标,并判断其是否超出图像范围,如果超出则跳过该检测结果或者2D边界框的右下角坐标的横坐标小于0或纵坐标小于0,则跳过该检测结果;如果当前输入数据是激光雷达,则判断当前3D边界框的中心点是否在规定范围内,如果不在则跳过该检测结果。
- 将2D边界框的坐标调整为左上角和右下角的形式,并将其调整为图像范围内的坐标。
- 创建一个字典
result_dict
,将当前检测结果的类别、旋转角度、2D边界框、3D边界框中心点、3D边界框尺寸和旋转角度、置信度等信息保存到该字典中。- 使用
kitti_result_line()
函数将当前检测结果的信息转换为Kitti数据集格式的一行,将其保存到result_lines
列表中。 值得注意的是,该函数中的kitti_result_line()
函数可以将检测结果转换为Kitti数据集格式中的一行,从而方便后续的结果输出。
def predict_kitti_to_anno(net,
example,
class_names,
center_limit_range=None,
lidar_input=False,
global_set=None):
net
:目标检测模型。example
:Kitti数据集格式的样本,包含了点云、图像、标定矩阵、3D边界框等信息。class_names
:目标检测模型所能识别的物体类别列表。center_limit_range
:可选的中心点限制范围,用于剔除超出范围的检测结果。lidar_input
:可选的布尔值,表示是否使用激光雷达点云作为输入。global_set
:可选的全局变量,用于存储一些全局配置参数。
batch_image_shape = example[9]
batch_imgidx = example[8]
# 每个点的x、y、z坐标和反射强度
# 都是形状为(1, 1, H, W)的四维张量,其中H和W分别为点云分割后的网格高度和宽度
pillar_x = example[0][:, :, 0].unsqueeze(0).unsqueeze(0)
pillar_y = example[0][:, :, 1].unsqueeze(0).unsqueeze(0)
pillar_z = example[0][:, :, 2].unsqueeze(0).unsqueeze(0)
pillar_i = example[0][:, :, 3].unsqueeze(0).unsqueeze(0)
num_points_per_pillar = example[1].float().unsqueeze(0)
# Find distance of x, y, and z from pillar center 查找x、y和z与支柱中心的距离
# assuming xyres_16.proto
# 点云坐标信息处理
coors_x = example[2][:, 3].float()
coors_y = example[2][:, 2].float()
x_sub = coors_x.unsqueeze(1) * 0.16 + 0.1
y_sub = coors_y.unsqueeze(1) * 0.16 + -39.9
ones = torch.ones([1, 100], dtype=torch.float32, device=pillar_x.device)
x_sub_shaped = torch.mm(x_sub, ones).unsqueeze(0).unsqueeze(0)
y_sub_shaped = torch.mm(y_sub, ones).unsqueeze(0).unsqueeze(0)
# 获取其它与目标检测相关的信息
num_points_for_a_pillar = pillar_x.size()[3] # 每个点云分割后的小块中实际存在的点数量
# 将神经网络输入中没有对应的点云小块的部分标记为无效数据
mask = get_paddings_indicator(num_points_per_pillar, num_points_for_a_pillar, axis=0)
mask = mask.permute(0, 2, 1)
mask = mask.unsqueeze(1)
mask = mask.type_as(pillar_x)
coors = example[2]
anchors = example[6]
anchors_mask = example[7]
anchors_mask = torch.as_tensor(anchors_mask, dtype=torch.uint8, device=pillar_x.device)
anchors_mask = anchors_mask.byte()
rect = example[3]
Trv2c = example[4]
P2 = example[5]
image_idx = example[8]
coors
:点云坐标信息,形状为(N, 4)
,其中N
为点的数量,每个点包括x、y、z坐标和所在的点云分割后的小块编号。anchors
:3D边界框的锚点信息,形状为(A, 7)
,其中A
为锚点的数量,每个锚点包括x、y、z坐标、边界框的宽度、高度、深度和旋转角度信息。anchors_mask
:锚点掩码信息,形状为(A,)
,其中每个元素为0或1,用于标记哪些锚点是有效的。rect
:点云坐标系对应的3D边界框的长宽高信息,形状为(N, 3)
,其中每个元素表示一个3D边界框的长宽高。Trv2c
:点云坐标系到相机坐标系的转换矩阵,形状为(4, 4)
。P2
:相机的投影矩阵,形状为(3, 4)
。image_idx
:输入的图像在数据集中的索引,一个标量。
for i, preds_dict in enumerate(predictions_dicts):
image_shape = batch_image_shape[i]
img_idx = preds_dict[5]
if preds_dict[0] is not None: # bbox list
box_2d_preds = preds_dict[0].detach().cpu().numpy() # bbox
box_preds = preds_dict[1].detach().cpu().numpy() # bbox3d_camera
scores = preds_dict[3].detach().cpu().numpy() # scores
box_preds_lidar = preds_dict[2].detach().cpu().numpy() # box3d_lidar
# write pred to file
label_preds = preds_dict[4].detach().cpu().numpy() # label_preds
anno = kitti.get_start_result_anno() # 获取KITTI数据集中目标检测任务的初始标注信息
num_example = 0
for box, box_lidar, bbox, score, label in zip(
box_preds, box_preds_lidar, box_2d_preds, scores,
label_preds):
if not lidar_input:
if bbox[0] > image_shape[1] or bbox[1] > image_shape[0]:
continue
if bbox[2] < 0 or bbox[3] < 0:
continue
# print(img_shape)
if center_limit_range is not None:
limit_range = np.array(center_limit_range)
if (np.any(box_lidar[:3] < limit_range[:3])
or np.any(box_lidar[:3] > limit_range[3:])):
continue
image_shape = [image_shape[0], image_shape[1]]
# 对预测结果中的2D边界框坐标信息进行了裁剪
bbox[2:] = np.minimum(bbox[2:], image_shape[::-1])
bbox[:2] = np.maximum(bbox[:2], [0, 0])
anno["name"].append(class_names[int(label)])
anno["truncated"].append(0.0)
anno["occluded"].append(0)
anno["alpha"].append(-np.arctan2(-box_lidar[1], box_lidar[0]) +
box[6])
anno["bbox"].append(bbox)
anno["dimensions"].append(box[3:6])
anno["location"].append(box[:3])
anno["rotation_y"].append(box[6])
# 确保目标的置信度得分不会重复
if global_set is not None:
for i in range(100000):
if score in global_set:
score -= 1 / 100000
else:
global_set.add(score)
break
anno["score"].append(score)
num_example += 1
# 将每个数据集中的标注信息(anno)存储到一个列表中(annos)
if num_example != 0:
anno = {n: np.stack(v) for n, v in anno.items()}
annos.append(anno)
else:
annos.append(kitti.empty_result_anno())
else:
annos.append(kitti.empty_result_anno())
# 为每个数据集的标注信息添加一个image_idx键值对,表示该标注信息对应的图像在数据集中的索引。
num_example = annos[-1]["name"].shape[0]
annos[-1]["image_idx"] = np.array(
[img_idx] * num_example, dtype=np.int64)
anno["name"]
表示预测目标的类别信息,
anno["truncated"]
、anno["occluded"]
表示目标的遮挡和截断情况,
anno["alpha"]
表示目标的方向角度,
anno["bbox"]
表示目标的2D边界框坐标信息,
anno["dimensions"]
表示目标的3D边界框的宽、高、深度信息,
anno["location"]
表示目标的3D边界框的位置信息,
anno["rotation_y"]
表示目标的3D边界框的旋转角度信息,
anno["score"]
表示目标的置信度得分。最后,该代码对全局的置信度得分集合
global_set
进行了更新,确保不会出现重复的得分值。
首先,会判断当前数据集中是否存在标注信息。
如果存在,则将每个键值对中的值以NumPy数组的形式进行堆叠,存储到一个字典(
anno
)中。具体来说,np.stack()
函数可以将相同形状的数组按照指定的轴进行堆叠,形成一个新的数组。在这里,该函数的作用是将标注信息字典中的所有值(例如,目标的类别、位置、大小、旋转角度、置信度得分等)沿着第0个轴进行堆叠,形成一个新的数组。这样处理之后,标注信息字典中的每个值都变成了一个形状为(num_example, *)
的数组,其中num_example
是数据集中的样本数。 接着,将整个标注信息字典(anno
)存储到一个列表(annos
)中。这个列表中的每个元素都代表一个数据集中的标注信息,对应着数据集中的一个样本。最后,如果当前数据集中不存在标注信息,则将一个空的标注信息字典(
kitti.empty_result_anno()
)存储到annos
列表中,以占位的形式表示该数据集中没有样本。
kitti.empty_result_anno()
是一个用于创建空标注信息字典的函数,它的返回值是一个包含指定键值对的字典,其中每个值都是一个空的NumPy数组。
def evaluate(config_path,
model_dir,
result_path=None,
predict_test=False,
ckpt_path=None,
ref_detfile=None,
pickle_result=True):
配置文件路径(
config_path
)模型目录路径(
model_dir
)结果文件路径(
result_path
)是否对测试集进行预测(
predict_test
)检查点文件路径(
ckpt_path
)参考检测文件路径(
ref_detfile
)是否将评估结果保存为Pickle文件(
pickle_result
)
# 将model_dir转换为绝对路径
model_dir = str(Path(model_dir).resolve())
# 结果文件命名
if predict_test:
result_name = 'predict_test'
else:
result_name = 'eval_results'
# 保存路径
if result_path is None:
model_dir = Path(model_dir)
result_path = model_dir / result_name
else:
result_path = pathlib.Path(result_path)
# 加载配置文件并创建配置对象
if isinstance(config_path, str):
config = pipeline_pb2.TrainEvalPipelineConfig()
with open(config_path, "r") as f:
proto_str = f.read()
text_format.Merge(proto_str, config)
else:
config = config_path
input_cfg = config.eval_input_reader
model_cfg = config.model.second
train_cfg = config.train_config
class_names = list(input_cfg.class_names)
center_limit_range = model_cfg.post_center_limit_range
输入读取器配置(
input_cfg
)模型配置(
model_cfg
)训练配置(
train_cfg
)从输入读取器配置中获取类别名称列表(
class_names
)从模型配置中获取后处理中心点限制范围参数(
center_limit_range
)
生成体素生成器
voxel_generator = voxel_builder.build(model_cfg.voxel_generator) # 体素生成器
bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]] # 获取点云范围
box_coder = box_coder_builder.build(model_cfg.box_coder) # 框编码器
target_assigner_cfg = model_cfg.target_assigner # 获取目标分配器的配置信息
target_assigner = target_assigner_builder.build(target_assigner_cfg,
bv_range, box_coder) # 目标分配器
net = second_builder.build(model_cfg, voxel_generator, target_assigner,
input_cfg.batch_size)
net.cuda()
# 将网络转换为混合精度(mixed precision)模式
if train_cfg.enable_mixed_precision:
net.half()
net.metrics_to_float()
net.convert_norm_to_float(net)
生成器是构成3D目标检测模型的必要组成部分,主要用于将点云数据转换为适合输入到深度学习网络中的数据。具体来说,这些生成器的作用如下:
- 体素生成器(voxel generator):将点云数据转换为三维体素表示,即将点云数据按照一定的规则(如3D网格)划分为小的体素(即三维像素),并将每个体素中的点云信息进行聚合,生成一个体素特征表示。这样做的目的是将点云数据从不规则的形式转换为规则的三维网格形式,方便深度学习网络处理。
- 框编码器(box coder):将3D目标的位置、大小等信息编码成网络预测的回归目标。具体来说,框编码器将3D目标的位置、大小等信息转换为一组数值,这些数值作为网络的回归目标,用于指导网络学习如何预测3D目标的位置、大小等信息。
- 目标分配器(target assigner):将3D目标与体素中的特征向量进行匹配,生成网络训练所需的标注数据。具体来说,目标分配器将3D目标的位置与体素的位置进行匹配,将匹配成功的体素的特征向量作为该3D目标的特征表示,并生成网络训练所需的标注数据(如分类标签、回归目标等)。 这些生成器的作用是将点云数据转换为深度学习网络能够处理的格式,并生成网络训练所需的标注数据,从而实现对3D目标的检测和定位。
首先使用
second_builder.build()
函数根据模型配置、体素生成器、目标分配器和输入配置构建3D目标检测网络(net
),并指定网络的批量大小(input_cfg.batch_size
)。 然后,使用net.cuda()
函数将网络移动到GPU上运行,以加速网络的训练和推断。
混合精度模式(mixed precision)是一种优化深度神经网络训练的技术,将网络中的部分操作使用低精度浮点数(如16位浮点数)来计算,以减少内存占用和计算时间,提高训练速度和效率。同时,混合精度模式还可以通过使用特殊的数值格式(如半精度浮点数)来提高存储效率,从而实现更大规模的网络训练和推断。
# 恢复已保存的模型参数
if ckpt_path is None:
# 尝试从已保存的最新检查点文件中恢复模型参数
torchplus.train.try_restore_latest_checkpoints(model_dir, [net])
else:
torchplus.train.restore(ckpt_path, net)
# 数据集以及数据加载器
eval_dataset = input_reader_builder.build(
input_cfg,
model_cfg,
training=False,
voxel_generator=voxel_generator,
target_assigner=target_assigner)
eval_dataloader = torch.utils.data.DataLoader(
eval_dataset,
batch_size=input_cfg.batch_size,
shuffle=False,
num_workers=input_cfg.num_workers,
pin_memory=False,
collate_fn=merge_second_batch)
# 根据训练配置选择浮点数精度类型
if train_cfg.enable_mixed_precision:
float_dtype = torch.float16
else:
float_dtype = torch.float32
net.eval() # 评估模式
result_path_step = result_path / f"step_{net.get_global_step()}"
result_path_step.mkdir(parents=True, exist_ok=True) # 保存每个全局步骤(net.get_global_step())的输出标签结果
t = time.time()
dt_annos = [] # 保存输出标签结果
global_set = None # 保存全局点云集合
print("Generate output labels...")
bar = ProgressBar() # 进度条
bar.start(len(eval_dataset) // input_cfg.batch_size + 1) # 进度条总长度
for example in iter(eval_dataloader):
# eval example [0: 'voxels', 1: 'num_points', 2: 'coordinates', 3: 'rect'
# 4: 'Trv2c', 5: 'P2', 6: 'anchors', 7: 'anchors_mask'
# 8: 'image_idx', 9: 'image_shape']
example = example_convert_to_torch(example, float_dtype)
example_tuple = list(example.values())
example_tuple[8] = torch.from_numpy(example_tuple[8])
example_tuple[9] = torch.from_numpy(example_tuple[9])
# 当前批次的样本数量是否为指定的批次大小
if (example_tuple[6].size()[0] != input_cfg.batch_size):
continue
if pickle_result:
# 对当前批次的样本进行目标检测,生成输出标签
dt_annos += predict_kitti_to_anno(
net, example_tuple, class_names, center_limit_range,
model_cfg.lidar_input, global_set)
else:
# 将输出标签保存到指定路径中
_predict_kitti_to_file(net, example, result_path_step, class_names,
center_limit_range, model_cfg.lidar_input)
bar.print_bar() # 打印进度条
# 每秒钟可以处理的样本数
sec_per_example = len(eval_dataset) / (time.time() - t)
print(f'generate label finished({sec_per_example:.2f}/s). start eval:')
print(f"avg forward time per example: {net.avg_forward_time:.3f}")
print(f"avg postprocess time per example: {net.avg_postprocess_time:.3f}")
如果该代码不是用于预测测试集,则首先从评估数据集中获取所有样本的真实标注信息
gt_annos
,并使用get_label_annos()
函数从模型输出的标签文件中获取所有样本的预测标注信息dt_annos
。 接着,该代码调用get_official_eval_result()
函数和get_coco_eval_result()
函数分别计算目标检测模型在KITTI评估指标和COCO评估指标下的性能表现,并输出评估结果。KITTI评估指标是用于评估自动驾驶环境下目标检测算法的性能的指标,包括平均精度(AP)、平均重召回率(AR)等;而COCO评估指标是用于评估通用物体检测算法的性能的指标,包括平均精度(AP)、平均重召回率(AR)等。 最后,如果指定了pickle_result
为True,则将评估结果保存到指定路径下的result.pkl
文件中。
def export_onnx(net, example, class_names, batch_image_shape,
center_limit_range=None, lidar_input=False, global_set=None):
将PyTorch模型转换为ONNX格式,并保存为文件。
net
:PyTorch模型实例;example
:示例数据,用于执行计算图并生成ONNX模型;class_names
:目标检测任务中的类别名称列表;batch_image_shape
:输入数据的形状,格式为(batch_size, channels, height, width);center_limit_range
:目标检测任务中物体中心点在Lidar坐标系下的取值范围;lidar_input
:是否使用Lidar点云数据作为输入;global_set
:其他设置参数。
def export_onnx(net, example, class_names, batch_image_shape,
center_limit_range=None, lidar_input=False, global_set=None):
pillar_x = example[0][:, :, 0].unsqueeze(0).unsqueeze(0)
pillar_y = example[0][:, :, 1].unsqueeze(0).unsqueeze(0)
pillar_z = example[0][:, :, 2].unsqueeze(0).unsqueeze(0)
pillar_i = example[0][:, :, 3].unsqueeze(0).unsqueeze(0)
num_points_per_pillar = example[1].float().unsqueeze(0)
# Find distance of x, y, and z from pillar center
# assuming xyres_16.proto
coors_x = example[2][:, 3].float()
coors_y = example[2][:, 2].float()
x_sub = coors_x.unsqueeze(1) * 0.16 + 0.1
y_sub = coors_y.unsqueeze(1) * 0.16 + -39.9
ones = torch.ones([1, 100], dtype=torch.float32, device=pillar_x.device)
x_sub_shaped = torch.mm(x_sub, ones).unsqueeze(0).unsqueeze(0)
y_sub_shaped = torch.mm(y_sub, ones).unsqueeze(0).unsqueeze(0)
num_points_for_a_pillar = pillar_x.size()[3]
mask = get_paddings_indicator(num_points_per_pillar, num_points_for_a_pillar, axis=0)
mask = mask.permute(0, 2, 1)
mask = mask.unsqueeze(1)
mask = mask.type_as(pillar_x)
coors = example[2]
print(pillar_x.size())
print(pillar_y.size())
print(pillar_z.size())
print(pillar_i.size())
print(num_points_per_pillar.size())
print(x_sub_shaped.size())
print(y_sub_shaped.size())
print(mask.size())
input_names = ["pillar_x", "pillar_y", "pillar_z", "pillar_i",
"num_points_per_pillar", "x_sub_shaped", "y_sub_shaped", "mask"]
# Wierd Convloution
# 每个Tensor对象的形状为(1, 1, 12000, 100),其中1表示batch_size,12000表示pillar的数量,100表示每个pillar中点的数量;
pillar_x = torch.ones([1, 1, 12000, 100], dtype=torch.float32, device=pillar_x.device)
pillar_y = torch.ones([1, 1, 12000, 100], dtype=torch.float32, device=pillar_x.device)
pillar_z = torch.ones([1, 1, 12000, 100], dtype=torch.float32, device=pillar_x.device)
pillar_i = torch.ones([1, 1, 12000, 100], dtype=torch.float32, device=pillar_x.device)
# 示每个pillar中的点数
num_points_per_pillar = torch.ones([1, 12000], dtype=torch.float32, device=pillar_x.device)
# 表示pillar中每个点的x、y坐标与pillar中心点x、y坐标的差值
x_sub_shaped = torch.ones([1, 1, 12000, 100], dtype=torch.float32, device=pillar_x.device)
y_sub_shaped = torch.ones([1, 1, 12000, 100], dtype=torch.float32, device=pillar_x.device)
# 表示每个pillar中点的掩码值
mask = torch.ones([1, 1, 12000, 100], dtype=torch.float32, device=pillar_x.device)
example1 = [pillar_x, pillar_y, pillar_z, pillar_i,
num_points_per_pillar, x_sub_shaped, y_sub_shaped, mask]
print('-------------- network readable visiual --------------')
torch.onnx.export(net, example1, "pfe.onnx", verbose=False, input_names=input_names)
print('pfe.onnx transfer success ...')
rpn_input = torch.ones([1, 64, 496, 432], dtype=torch.float32, device=pillar_x.device)
torch.onnx.export(net.rpn, rpn_input, "rpn.onnx", verbose=False)
print('rpn.onnx transfer success ...')
return 0
def onnx_model_generate(config_path,
model_dir,
result_path=None,
predict_test=False,
ckpt_path=None
):
生成ONNX格式的模型。具体来说,该函数接收五个参数:
config_path
:配置文件路径,指定了模型的参数配置;model_dir
:模型保存目录,指定了模型的保存路径;result_path
:结果保存路径,指定了模型在测试集上的预测结果保存路径;predict_test
:是否在测试集上进行预测;ckpt_path
:模型权重文件路径,指定了模型权重的保存路径。
def onnx_model_generate(config_path,
model_dir,
result_path=None,
predict_test=False,
ckpt_path=None
):
model_dir = pathlib.Path(model_dir)
if predict_test:
result_name = 'predict_test'
else:
result_name = 'eval_results'
if result_path is None:
result_path = model_dir / result_name
else:
result_path = pathlib.Path(result_path)
config = pipeline_pb2.TrainEvalPipelineConfig()
with open(config_path, "r") as f:
proto_str = f.read()
text_format.Merge(proto_str, config)
input_cfg = config.eval_input_reader
model_cfg = config.model.second
train_cfg = config.train_config
class_names = list(input_cfg.class_names)
center_limit_range = model_cfg.post_center_limit_range
##########################
## Build Voxel Generator
##########################
voxel_generator = voxel_builder.build(model_cfg.voxel_generator)
bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]]
box_coder = box_coder_builder.build(model_cfg.box_coder)
target_assigner_cfg = model_cfg.target_assigner
target_assigner = target_assigner_builder.build(target_assigner_cfg,
bv_range, box_coder)
net = second_builder.build(model_cfg, voxel_generator, target_assigner, 1)
net.cuda()
if train_cfg.enable_mixed_precision:
net.half()
net.metrics_to_float()
net.convert_norm_to_float(net)
if ckpt_path is None:
torchplus.train.try_restore_latest_checkpoints(model_dir, [net])
else:
torchplus.train.restore(ckpt_path, net)
eval_dataset = input_reader_builder.build(
input_cfg,
model_cfg,
training=False,
voxel_generator=voxel_generator,
target_assigner=target_assigner)
eval_dataloader = torch.utils.data.DataLoader(
eval_dataset,
batch_size=1,
shuffle=False,
num_workers=1,
pin_memory=False,
collate_fn=merge_second_batch)
if train_cfg.enable_mixed_precision:
float_dtype = torch.float16
else:
float_dtype = torch.float32
net.eval()
result_path_step = result_path / f"step_{net.get_global_step()}"
result_path_step.mkdir(parents=True, exist_ok=True)
dt_annos = []
global_set = None
print("Generate output labels...")
bar = ProgressBar()
bar.start(len(eval_dataset) // input_cfg.batch_size + 1)
for example in iter(eval_dataloader):
example = example_convert_to_torch(example, float_dtype)
example_tuple = list(example.values())
batch_image_shape = example_tuple[8]
example_tuple[8] = torch.from_numpy(example_tuple[8])
example_tuple[9] = torch.from_numpy(example_tuple[9])
dt_annos = export_onnx(
net, example_tuple, class_names, batch_image_shape, center_limit_range,
model_cfg.lidar_input, global_set)
return 0
bar.print_bar()
训练文件就暂时看到这,太多了,网络结构重开一篇吧。