网站域名和密码,网页设计和网站制作,企业网站建设哪里做网站好,京东物流网站1.模型训练
训练流程#xff1a;train.py中的training函数 这段代码实现了一个 3D 高斯模型的训练循环#xff0c;旨在通过逐步优化模型参数#xff0c;使其能够精确地渲染特定场景。以下是代码的详细解析#xff1a;
def training(dataset, opt, pipe, testing_iteratio…
1.模型训练
训练流程train.py中的training函数 这段代码实现了一个 3D 高斯模型的训练循环旨在通过逐步优化模型参数使其能够精确地渲染特定场景。以下是代码的详细解析
def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoint_iterations, checkpoint, debug_from):first_iter 0# 初始化高斯模型用于表示场景中的每个点的3D高斯分布gaussians GaussianModel(dataset.sh_degree)# 初始化场景对象加载数据集和对应的相机参数scene Scene(dataset, gaussians)# 为高斯模型参数设置优化器和学习率调度器gaussians.training_setup(opt)# 如果提供了checkpoint则从checkpoint加载模型参数并恢复训练进度if checkpoint:(model_params, first_iter) torch.load(checkpoint)gaussians.restore(model_params, opt)# 设置背景颜色白色或黑色取决于数据集要求bg_color [1, 1, 1] if dataset.white_background else [0, 0, 0]background torch.tensor(bg_color, dtypetorch.float32, devicecuda)# 创建CUDA事件用于计时iter_start torch.cuda.Event(enable_timingTrue)iter_end torch.cuda.Event(enable_timingTrue)viewpoint_stack Noneema_loss_for_log 0.0# 使用tqdm库创建进度条追踪训练进度progress_bar tqdm(range(first_iter, opt.iterations), descTraining progress)first_iter 1for iteration in range(first_iter, opt.iterations 1):# 记录迭代开始时间iter_start.record()# 根据当前迭代次数更新学习率gaussians.update_learning_rate(iteration)# 每1000次迭代提升球谐函数的次数以改进模型复杂度if iteration % 1000 0:gaussians.oneupSHdegree()# 随机选择一个训练用的相机视角if not viewpoint_stack:viewpoint_stack scene.getTrainCameras().copy()viewpoint_cam viewpoint_stack.pop(randint(0, len(viewpoint_stack)-1))# 如果达到调试起始点启用调试模式if (iteration - 1) debug_from:pipe.debug True# 根据设置决定是否使用随机背景颜色bg torch.rand((3), devicecuda) if opt.random_background else background# 渲染当前视角的图像render_pkg render(viewpoint_cam, gaussians, pipe, bg)image, viewspace_point_tensor, visibility_filter, radii render_pkg[render], render_pkg[viewspace_points], render_pkg[visibility_filter], render_pkg[radii]# 计算渲染图像与真实图像之间的损失gt_image viewpoint_cam.original_image.cuda()Ll1 l1_loss(image, gt_image)loss (1.0 - opt.lambda_dssim) * Ll1 opt.lambda_dssim * (1.0 - ssim(image, gt_image))loss.backward()# 记录迭代结束时间iter_end.record()with torch.no_grad():# 更新进度条和损失显示ema_loss_for_log 0.4 * loss.item() 0.6 * ema_loss_for_logif iteration % 10 0:progress_bar.set_postfix({Loss: f{ema_loss_for_log:.{7}f}})progress_bar.update(10)if iteration opt.iterations:progress_bar.close()# 定期记录训练数据并保存模型training_report(tb_writer, iteration, Ll1, loss, l1_loss, iter_start.elapsed_time(iter_end), testing_iterations, scene, render, (pipe, background))if iteration in saving_iterations:print(\n[ITER {}] Saving Gaussians.format(iteration))scene.save(iteration)# 在指定迭代区间内对3D高斯模型进行增密和修剪if iteration lt; opt.densify_until_iter:gaussians.max_radii2D[visibility_filter] torch.max(gaussians.max_radii2D[visibility_filter], radii[visibility_filter])gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter)if iteration gt; opt.densify_from_iter and iteration % opt.densification_interval 0:size_threshold 20 if iteration gt; opt.opacity_reset_interval else Nonegaussians.densify_and_prune(opt.densify_grad_threshold, 0.005, scene.cameras_extent, size_threshold)if iteration % opt.opacity_reset_interval 0 or (dataset.white_background and iteration opt.densify_from_iter):gaussians.reset_opacity()# 执行优化器的一步并准备下一次迭代if iteration lt; opt.iterations:gaussians.optimizer.step()gaussians.optimizer.zero_grad(set_to_noneTrue)# 定期保存checkpointif iteration in checkpoint_iterations:print(\n[ITER {}] Saving Checkpoint.format(iteration))torch.save((gaussians.capture(), iteration), scene.model_path /chkpnt str(iteration) .pth)函数解析 training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoint_iterations, checkpoint, debug_from) 该函数的作用是加载数据、设置场景和高斯模型的初始化参数通过迭代训练模型调整模型参数以减小渲染图像与真实图像之间的误差。 参数 dataset: 数据集对象包含图像和相机参数。opt: 配置对象包含训练过程的超参数如学习率和迭代次数。pipe: 渲染管道用于生成渲染图像。testing_iterations: 在这些迭代次数中进行测试。saving_iterations: 在这些迭代次数中保存高斯模型。checkpoint_iterations: 在这些迭代次数中保存 checkpoint。checkpoint: 如果提供了 checkpoint 文件则从中恢复模型状态。debug_from: 设定调试起始点达到该迭代后启用调试模式。 整个函数旨在通过高斯模型增密、学习率调整、定期保存 checkpoint 和渲染优化训练一个能够表示 3D 场景的高斯模型。
2.数据加载
scene/__init__.py中的Scene类的__init__函数
2.1Class Scene 这段代码定义了 Scene 类用于管理和加载 3D 场景的参数、模型和相机信息并支持不同分辨率的相机数据。Scene 类结合了高斯模型和数据集处理逻辑尤其适用于从 COLMAP 或 Blender 数据集中加载相机和场景信息以便进行 3D 表示和训练。 class Scene:Scene 类用于管理场景的3D模型包括相机参数、点云数据和高斯模型的初始化和加载def __init__(self, args: ModelParams, gaussians: GaussianModel, load_iterationNone, shuffleTrue, resolution_scales[1.0]):初始化场景对象:param args: 包含模型路径和源路径等模型参数:param gaussians: 高斯模型对象用于场景点的3D表示:param load_iteration: 指定加载模型的迭代次数如果为-1则自动寻找最大迭代次数:param shuffle: 是否在训练前打乱相机列表:param resolution_scales: 分辨率比例列表用于处理不同分辨率的相机self.model_path args.model_path # 模型文件保存路径self.loaded_iter None # 已加载的迭代次数self.gaussians gaussians # 高斯模型对象# 检查并加载已有的训练模型if load_iteration:if load_iteration -1:self.loaded_iter searchForMaxIteration(os.path.join(self.model_path, point_cloud))else:self.loaded_iter load_iterationprint(fLoading trained model at iteration {self.loaded_iter})self.train_cameras {} # 用于训练的相机参数self.test_cameras {} # 用于测试的相机参数# 根据数据集类型COLMAP或Blender加载场景信息if os.path.exists(os.path.join(args.source_path, sparse)):scene_info sceneLoadTypeCallbacks[Colmap](args.source_path, args.images, args.eval)elif os.path.exists(os.path.join(args.source_path, transforms_train.json)):print(Found transforms_train.json file, assuming Blender data set!)scene_info sceneLoadTypeCallbacks[Blender](args.source_path, args.white_background, args.eval)else:assert False, Could not recognize scene type!# 如果是初次训练初始化3D高斯模型否则加载已有模型if self.loaded_iter:self.gaussians.load_ply(os.path.join(self.model_path, point_cloud, iteration_ str(self.loaded_iter), point_cloud.ply))else:self.gaussians.create_from_pcd(scene_info.point_cloud, self.cameras_extent)# 根据resolution_scales加载不同分辨率的训练和测试位姿for resolution_scale in resolution_scales:print(Loading Training Cameras)self.train_cameras[resolution_scale] cameraList_from_camInfos(scene_info.train_cameras, resolution_scale, args)print(Loading Test Cameras)self.test_cameras[resolution_scale] cameraList_from_camInfos(scene_info.test_cameras, resolution_scale, args)def save(self, iteration):保存当前迭代下的3D高斯模型点云。:param iteration: 当前的迭代次数。point_cloud_path os.path.join(self.model_path, fpoint_cloud/iteration_{iteration})self.gaussians.save_ply(os.path.join(point_cloud_path, point_cloud.ply))def getTrainCameras(self, scale1.0):获取指定分辨率比例的训练相机列表:param scale: 分辨率比例:return: 指定分辨率比例的训练相机列表return self.train_cameras[scale]sceneLoadTypeCallbacks {Colmap: readColmapSceneInfo,Blender : readNerfSyntheticInfo
}类属性和方法解析 1. __init__(self, args: ModelParams, gaussians: GaussianModel, load_iterationNone, shuffleTrue, resolution_scales[1.0]) 构造函数 __init__ 用于初始化 Scene 对象加载数据集和初始化 3D 高斯模型并为训练和测试分辨率管理相机参数。 参数 args包含模型路径和源路径等配置参数。gaussians高斯模型对象用于管理场景中的 3D 高斯分布。load_iteration指定加载的迭代次数若为 -1自动搜索最后保存的模型。shuffle是否对相机列表进行随机排序。resolution_scales用于管理不同分辨率下的相机数据。 主要步骤 加载已有模型如果 load_iteration 不为 None加载指定迭代次数的模型若 load_iteration -1加载最新迭代的模型。场景信息加载通过回调函数 sceneLoadTypeCallbacks从 COLMAP 或 Blender 数据集中加载场景信息。初始化高斯模型如果没有加载已有的模型调用 create_from_pcd 方法使用从数据集中加载的点云初始化高斯模型。分辨率管理根据 resolution_scales 参数加载不同分辨率的训练和测试相机列表分别存储在 self.train_cameras 和 self.test_cameras 字典中。 2. save(self, iteration) 该方法用于在指定的 iteration 下保存当前 3D 高斯模型的点云数据。 参数 iteration当前的迭代次数。 实现 根据 iteration 创建路径并调用 self.gaussians.save_ply 保存点云数据。 3. getTrainCameras(self, scale1.0) 该方法用于获取指定分辨率比例的训练相机列表。 参数 scale分辨率比例默认为 1.0用来返回该分辨率对应的训练相机列表。 返回值 以 scale 为键从 self.train_cameras 中获取对应的相机列表。 4. sceneLoadTypeCallbacks sceneLoadTypeCallbacks 是一个字典用于匹配和回调不同数据集的加载函数 如果数据集是 COLMAP 类型调用 readColmapSceneInfo。如果是 Blender 数据集调用 readNerfSyntheticInfo。 示例流程 初始化 Scene 类根据数据集路径自动检测数据类型并调用合适的回调函数加载相机和场景信息。 分辨率和相机管理 支持不同分辨率的训练和测试相机列表。用户可以通过 getTrainCameras(scale) 获取指定分辨率下的训练相机。 保存 在指定的迭代次数 iteration 下调用 save() 将 3D 高斯模型点云数据保存为 .ply 文件。 代码优点 模块化和扩展性支持多种数据集格式并允许通过回调函数轻松扩展至其他格式。分辨率管理灵活加载和管理不同分辨率的相机列表以便在训练和测试中使用不同精度的相机数据 2.2 readColmapSceneInfo 和 readColmapCameras
这段代码提供了两个主要函数 readColmapSceneInfo 和 readColmapCameras用于读取和处理 COLMAP 生成的相机参数和场景信息并将其格式化为易于使用的数据结构。这些函数为后续 3D 重建和 NeRF 训练准备相机参数、点云数据及其他场景归一化信息。
def readColmapSceneInfo(path, images, eval, llffhold8):# 尝试读取COLMAP处理结果中的二进制相机外参和内参文件try:cameras_extrinsic_file os.path.join(path, sparse/0, images.bin)cameras_intrinsic_file os.path.join(path, sparse/0, cameras.bin)cam_extrinsics read_extrinsics_binary(cameras_extrinsic_file)cam_intrinsics read_intrinsics_binary(cameras_intrinsic_file)except:# 如果二进制文件读取失败尝试读取文本格式的相机外参和内参文件cameras_extrinsic_file os.path.join(path, sparse/0, images.txt)cameras_intrinsic_file os.path.join(path, sparse/0, cameras.txt)cam_extrinsics read_extrinsics_text(cameras_extrinsic_file)cam_intrinsics read_intrinsics_text(cameras_intrinsic_file)# 定义存放图片的目录如果未指定则默认为imagesreading_dir images if images is None else images# 读取并处理相机参数转换为内部使用的格式cam_infos_unsorted readColmapCameras(cam_extrinsicscam_extrinsics, cam_intrinsicscam_intrinsics, images_folderos.path.join(path, reading_dir))# 根据图片名称对相机信息进行排序以保证顺序一致性cam_infos sorted(cam_infos_unsorted.copy(), keylambda x: x.image_name)# 根据是否为评估模式eval将相机分为训练集和测试集# 如果为评估模式根据llffhold参数通常用于LLFF数据集间隔选择测试相机if eval:train_cam_infos [c for idx, c in enumerate(cam_infos) if idx % llffhold ! 0]test_cam_infos [c for idx, c in enumerate(cam_infos) if idx % llffhold 0]else:# 如果不是评估模式所有相机均为训练相机测试相机列表为空train_cam_infos cam_infostest_cam_infos []# 计算场景归一化参数这是为了处理不同尺寸和位置的场景使模型训练更稳定nerf_normalization getNerfppNorm(train_cam_infos)# 尝试读取点云数据优先从PLY文件读取如果不存在则尝试从BIN或TXT文件转换并保存为PLY格式ply_path os.path.join(path, sparse/0/points3D.ply)bin_path os.path.join(path, sparse/0/points3D.bin)txt_path os.path.join(path, sparse/0/points3D.txt)if not os.path.exists(ply_path):print(Converting point3d.bin to .ply, will happen only the first time you open the scene.)try:xyz, rgb, _ read_points3D_binary(bin_path)except:xyz, rgb, _ read_points3D_text(txt_path)storePly(ply_path, xyz, rgb)try:pcd fetchPly(ply_path)except:pcd None# 组装场景信息包括点云、训练用相机、测试用相机、场景归一化参数和点云文件路径scene_info SceneInfo(point_cloudpcd,train_camerastrain_cam_infos,test_camerastest_cam_infos,nerf_normalizationnerf_normalization,ply_pathply_path)return scene_infodef readColmapCameras(cam_extrinsics, cam_intrinsics, images_folder):cam_infos [] # 初始化用于存储相机信息的列表# 遍历所有相机的外参for idx, key in enumerate(cam_extrinsics):# 动态显示读取相机信息的进度sys.stdout.write(\r)sys.stdout.write(Reading camera {}/{}.format(idx1, len(cam_extrinsics)))sys.stdout.flush()# 获取当前相机的外参和内参extr cam_extrinsics[key] # 当前相机的外参intr cam_intrinsics[extr.camera_id] # 根据外参中的camera_id找到对应的内参height intr.height # 相机图片的高度width intr.width # 相机图片的宽度uid intr.id # 相机的唯一标识符# 将四元数表示的旋转转换为旋转矩阵RR np.transpose(qvec2rotmat(extr.qvec))# 外参中的平移向量T np.array(extr.tvec)# 根据相机内参模型计算视场角FoVif intr.model SIMPLE_PINHOLE:# 如果是简单针孔模型只有一个焦距参数focal_length_x intr.params[0]FovY focal2fov(focal_length_x, height) # 计算垂直方向的视场角FovX focal2fov(focal_length_x, width) # 计算水平方向的视场角elif intr.model PINHOLE:# 如果是针孔模型有两个焦距参数focal_length_x intr.params[0]focal_length_y intr.params[1]FovY focal2fov(focal_length_y, height) # 使用y方向的焦距计算垂直视场角FovX focal2fov(focal_length_x, width) # 使用x方向的焦距计算水平视场角else:# 如果不是以上两种模型抛出错误assert False, Colmap camera model not handled: only undistorted datasets (PINHOLE or SIMPLE_PINHOLE cameras) supported!# 构建图片的完整路径image_path os.path.join(images_folder, os.path.basename(extr.name))image_name os.path.basename(image_path).split(.)[0] # 提取图片名称不包含扩展名# 使用PIL库打开图片文件image Image.open(image_path)# 创建并存储相机信息cam_info CameraInfo(uiduid, RR, TT, FovYFovY, FovXFovX, imageimage,image_pathimage_path, image_nameimage_name, widthwidth, heightheight)cam_infos.append(cam_info)# 在读取完所有相机信息后换行sys.stdout.write(\n)# 返回整理好的相机信息列表return cam_infos
函数解析
readColmapSceneInfo(path, images, eval, llffhold8) 该函数负责读取和解析 COLMAP 数据集的场景信息包括相机参数、点云和归一化参数等。主要步骤如下 相机参数文件读取 首先尝试读取二进制格式的相机外参 (images.bin) 和内参 (cameras.bin) 文件若读取失败则尝试加载文本格式的文件images.txt 和 cameras.txt。外参和内参分别通过 read_extrinsics_binary 和 read_intrinsics_binary 等函数加载。 相机信息解析 调用 readColmapCameras 函数解析外参和内参文件中的相机数据并返回相机列表 cam_infos。 训练和测试集划分 根据 eval 参数决定是否将数据集划分为训练集和测试集。如果为评估模式 (evalTrue)会根据 llffhold 参数通常用于 LLFF 数据集进行间隔划分得到 train_cam_infos 和 test_cam_infos。 场景归一化参数计算 调用 getNerfppNorm 函数对训练集中的相机位姿进行归一化确保不同场景的尺度和位置对齐以便更稳定的训练。 点云文件读取与转换 点云文件优先从 points3D.ply 文件中读取若不存在则尝试从二进制或文本格式的 points3D 文件中读取并转换为 .ply 文件存储。最终通过 fetchPly 函数读取 .ply 格式的点云数据。 返回值 将 SceneInfo 对象返回其中包含点云数据、训练相机、测试相机、归一化参数和点云文件路径。 readColmapCameras(cam_extrinsics, cam_intrinsics, images_folder) 该函数用于将 cam_extrinsics 和 cam_intrinsics 中的相机数据解析成 CameraInfo 对象的列表并处理每个相机的旋转、平移、视场角等信息。 主要步骤 迭代解析相机外参遍历 cam_extrinsics 字典中的每个相机获取旋转矩阵 R 和平移向量 T。视场角计算根据相机模型SIMPLE_PINHOLE 或 PINHOLE计算垂直和水平方向的视场角。图像路径与信息结合图像文件夹路径 images_folder生成图像路径和图像名称并打开图像文件。创建并存储相机信息将解析到的相机信息存储在 CameraInfo 对象中并添加到 cam_infos 列表。 返回值 返回 cam_infos包含所有解析过的相机信息列表。 代码优点
自动文件读取和兼容性自动适配二进制和文本格式的文件读取使得代码更加兼容不同格式的 COLMAP 输出。训练/测试集划分灵活支持按一定间隔划分训练和测试集适合用于评估模式的数据集处理。相机视场角处理根据相机模型选择不同的视场角计算方式适应多种相机模型。 代码潜在改进异常处理在图像加载和文件读取时增加更多异常处理以应对图像丢失或格式不匹配的情况。多线程优化如果相机数量较多可使用多线程或并行处理以加快相机参数的读取速度。
