# # Copyright (C) 2023, Inria # GRAPHDECO research group, https://team.inria.fr/graphdeco # All rights reserved. # # This software is free for non-commercial, research and evaluation use # under the terms of the LICENSE.md file. # # For inquiries contact george.drettakis@inria.fr # import os import torch from random import randint from utils.loss_utils import l1_loss, ssim from gaussian_renderer import render, network_gui import sys from scene import Scene, GaussianModel, eghuman_dataset, collate_fn, Eghuman from scene.gaussian_nograd_model import GaussianModel_nograd from scene.FLAME import FLAME from utils.general_utils import safe_state import uuid from tqdm import tqdm import numpy as np from utils.image_utils import psnr from argparse import ArgumentParser, Namespace from arguments import ModelParams, PipelineParams, OptimizationParams from utils.sh_utils import RGB2SH from torch.utils.data import DataLoader import torch.nn.functional as F # from torch_geometric.data import Data # from torch_geometric.transforms import SamplePoints try: from torch.utils.tensorboard import SummaryWriter TENSORBOARD_FOUND = False except ImportError: TENSORBOARD_FOUND = False os.environ["CUDA_VISIBLE_DEVICES"] = "0" os.environ['CUDA_LAUNCH_BLOCKING'] = '1' def encoding_position(xyz_sample): # n_log2 = int(np.log2(tex_h)) # # 4, 8, ..., tex_h # channel_list = [2**i for i in range(2, n_log2 + 1)] L = 10 # +X right, +Y down, +Z forward # [#planes, tex_h, tex_w, 1] # d_eoncode_dict = {} # for _, tmp_res in enumerate(channel_list): posison_encoding = [] for i in range(L): posison_encoding.append(torch.sin(2**i*np.pi*xyz_sample)) posison_encoding.append(torch.cos(2**i*np.pi*xyz_sample)) posison_encoding.append(torch.sin(2**L*np.pi*xyz_sample)) posison_encoding = torch.concat(posison_encoding, 1) return posison_encoding def requires_grad(model, flag=True): for p in model.parameters(): p.requires_grad = flag def gs_requires_grad(gs, flag): for i in gs: for p in i['params']: p.requires_grad = flag # gs._xyz_scale = gs._xyz_scale.detach() # gs._align_R = gs._align_R.detach() # gs._align_T = gs._align_T.detach() # gs._features_dc = gs._features_dc.detach() # gs._features_rest = gs._features_rest.detach() # gs._opacity = gs._opacity.detach() # gs._scaling = gs._scaling.detach() # gs ._rotation = gs._rotation.detach() def training(train_loader, dataset, opt, pipe, testing_iterations, saving_iterations, checkpoint_iterations, checkpoint, debug_from): # gaussians = GaussianModel(dataset.sh_degree, dataset.data_device, dataset.xyz_grad, dataset.align, dataset.free_flame) # scene = Scene(dataset, gaussians) # flame = FLAME(dataset, opt) # flame.set_up() eghuman = Eghuman(opt, dataset.data_device, dataset.use_sh, dataset.sh_degree, dataset.checkpoint_path) eghuman.set_up() # gs_nograd = GaussianModel_nograd() epoch = 0 for step ,data_train in enumerate(train_loader): gaussians = data_train[0]['gs'] scene = data_train[0]['sc'] dataset = data_train[0]['args'] first_iter = 0 tb_writer = prepare_output_and_logger(dataset) gaussians.training_setup(opt) 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, dtype=torch.float32, device=dataset.data_device) iter_start = torch.cuda.Event(enable_timing = True) iter_end = torch.cuda.Event(enable_timing = True) viewpoint_stack = None ema_loss_for_log = 0.0 progress_bar = tqdm(range(first_iter, opt.iterations), desc="align flame progress") first_iter += 1 requires_grad(eghuman.tri_encoder, False) requires_grad(eghuman.gs_decoder, False) # requires_grad(flame.E_flame, True) for iteration in range(first_iter, opt.iterations + 1): if network_gui.conn == None: network_gui.try_connect() while network_gui.conn != None: try: net_image_bytes = None custom_cam, do_training, pipe.convert_SHs_python, pipe.compute_cov3D_python, keep_alive, scaling_modifer = network_gui.receive() if custom_cam != None: net_image = render(custom_cam, gaussians, pipe, background, scaling_modifer)["render"] net_image_bytes = memoryview((torch.clamp(net_image, min=0, max=1.0) * 255).byte().permute(1, 2, 0).contiguous().cpu().numpy()) network_gui.send(net_image_bytes, dataset.source_path) if do_training and ((iteration < int(opt.iterations)) or not keep_alive): break except Exception as e: network_gui.conn = None iter_start.record() gaussians.update_learning_rate(iteration) # Every 1000 its we increase the levels of SH up to a maximum degree if iteration % 1000 == 0: gaussians.oneupSHdegree() # Pick a random Camera if not viewpoint_stack: viewpoint_stack = scene.getTrainCameras().copy() if dataset.pick_camera: viewpoint_cam = viewpoint_stack[scene.view_index] else: viewpoint_cam = viewpoint_stack[60] # viewpoint_cam = viewpoint_stack.pop(randint(0, len(viewpoint_stack)-1)) if (iteration - 1) == debug_from: pipe.debug = True # flame_input = F.interpolate(gt_image.unsqueeze(0), (224, 224)) # codedict = flame.E_flame(flame_input, opt.n_shape, opt.n_exp, opt.n_pose) # verts = flame(shape_params=codedict['shape'], expression_params=codedict['exp'], pose_params=codedict['pose']) # gaussians._origin_xyz = verts.squeeze() # Render if not dataset.xyz_grad and dataset.align: gaussians.update_gs(gaussians.get_xyz) if dataset.use_sh: render_pkg = render(viewpoint_cam, gaussians, pipe, background) else: render_pkg = render(viewpoint_cam, gaussians, pipe, background, override_color = gaussians._features_dc) image, viewspace_point_tensor, visibility_filter, radii = render_pkg["render"], render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg["radii"] # print(torch.mean(radii.float())) # Loss gt_image = (viewpoint_cam.original_image).to(dataset.data_device) Ll1 = l1_loss(image, gt_image) loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim(image, gt_image)) loss.backward(retain_graph=True) iter_end.record() with torch.no_grad(): # Progress bar ema_loss_for_log = 0.4 * loss.item() + 0.6 * ema_loss_for_log if 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() # Log training_report(tb_writer, iteration, Ll1, loss, l1_loss, iter_start.elapsed_time(iter_end), testing_iterations, scene, render, (pipe, background)) # Densification # if iteration < opt.densify_until_iter: # # Keep track of max radii in image-space for pruning # 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 > opt.densify_from_iter and iteration % opt.densification_interval == 0: # size_threshold = 20 if iteration > opt.opacity_reset_interval else None # gaussians.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() # Optimizer step if iteration < opt.iterations: # flame.opter.step() gaussians.optimizer.step() gaussians.optimizer.zero_grad(set_to_none = True) # flame.opter.zero_grad(set_to_none = True) if (iteration in saving_iterations): #save print("\n[ITER {}] Saving FLAME".format(iteration)) scene.save(iteration, 'ply') if dataset.type_3d == 'obj': scene.save(iteration, dataset.type_3d) requires_grad(eghuman.tri_encoder, True) requires_grad(eghuman.gs_decoder, True) # requires_grad(flame.E_flame, False) gs_requires_grad(gaussians.l, False) #unsample Flame # sparse_flame = Data(pos=gaussians.get_xyz(), face=scene.face) # dense_flame = SamplePoints(50230)(sparse_flame) # gaussians._xyz = dense_flame #encoding tri_plane position # x_onehot = torch.tensor([[0., 0., 1.]]).repeat(16, 1).to(dataset.data_device) # y_onehot = torch.tensor([[0., 1., 0.]]).repeat(16, 1).to(dataset.data_device) # z_onehot = torch.tensor([[1., 0., 0.]]).repeat(16, 1).to(dataset.data_device) # # xyz_onehot = torch.concat((x_onehot, y_onehot, z_onehot), 1) x_min, x_max = torch.min(gaussians._xyz[:, 0]), torch.max(gaussians._xyz[:, 0]) y_min, y_max = torch.min(-gaussians._xyz[:, 1]), torch.max(-gaussians._xyz[:, 1]) z_min, z_max = torch.min(gaussians._xyz[:, 2]), torch.max(gaussians._xyz[:, 2]) # x_lin = torch.linspace(x_min.item(), x_max.item(), 16).to(dataset.data_device) # y_lin = torch.linspace(y_min.item(), y_max.item(), 16).to(dataset.data_device) # z_lin = torch.linspace(z_min.item(), z_max.item(), 16).to(dataset.data_device) # x_sample = torch.concat((x_onehot, F.normalize(x_lin.unsqueeze(1), 0)), 1) # y_sample = torch.concat((y_onehot, F.normalize(y_lin.unsqueeze(1), 0)), 1) # z_sample = torch.concat((z_onehot, F.normalize(z_lin.unsqueeze(1), 0)), 1) # xyz_sample = torch.concat((x_sample, y_sample, z_sample), 0) # xyz_embedding = encoding_position(xyz_sample) #process tri_plane_encoder viewpoint_stack = scene.getTrainCameras().copy() if dataset.pick_camera: input_viewpoint = viewpoint_stack.pop(scene.view_index) else: input_viewpoint = viewpoint_stack.pop(60) input_img = input_viewpoint.original_image.to(dataset.data_device).unsqueeze(0)#.repeat(3, 1, 1, 1) iter_start_net = torch.cuda.Event(enable_timing = True) iter_end_net = torch.cuda.Event(enable_timing = True) training_bar = tqdm(range(opt.iterations_human), desc="train eghuman progress") # gaussians.detach_para() features_dc = gaussians._features_dc.detach() features_rest = gaussians._features_rest.detach() for net_itr in range(opt.iterations_human): if not viewpoint_stack: viewpoint_stack = scene.getTrainCameras().copy() gt_viewpoint = viewpoint_stack.pop(randint(0, len(viewpoint_stack)-1)) gt_img = gt_viewpoint.original_image.to(dataset.data_device) # input_img_low = F.interpolate(input_img, (opt.img_size // 2, opt.img_size // 2)) iter_start_net.record() eghuman.opt.zero_grad() tri_feature = eghuman.tri_encoder(input_img) # point_f = eghuman.tri_grid_sample(tri_feature, gaussians.get_xyz, x_min, x_max, y_min, y_max, z_min, z_max) point_f3d = eghuman.tri_plane_sample(tri_feature, gaussians.get_xyz, x_min, x_max, y_min, y_max, z_min, z_max) feature_map, _ = eghuman.feature_net(input_img) point_uv = eghuman.uv_compute(gaussians.get_xyz, gt_viewpoint, dataset.data_device) point_fc2d = eghuman.uv_sample(input_img, feature_map, point_uv) if not dataset.use_sh: point_fc3d = torch.concat((features_dc.unsqueeze(0).repeat(1, 1, 3, 1), point_f3d), 3) gs_para = eghuman.gs_decoder(point_fc3d, point_fc2d, dataset.use_sh) else: num_points, _, _= gaussians._features_dc.shape point_fc3d = torch.concat((point_f3d, features_dc.unsqueeze(0).repeat(1, 1, 3, 1), torch.reshape(features_rest.unsqueeze(0), (1, num_points, -1)).unsqueeze(2).repeat(1, 1, 3, 1)), 3) gs_para = eghuman.gs_decoder(point_fc3d, point_fc2d, dataset.use_sh) gaussians._features_dc = gs_para['f_dc'] gaussians._opacity = gs_para['opacity'] # gaussians._scaling = gs_para['scaling'] # gaussians._rotation = gs_para['rotation'] if dataset.use_sh: gaussians._features_rest = gs_para['f_rest'] if dataset.use_sh: render_pkg = render(gt_viewpoint, gaussians, pipe, background) else: render_pkg = render(gt_viewpoint, gaussians, pipe, background, override_color = gaussians._features_dc) image_gs = render_pkg["render"] Ll1_net = l1_loss(image_gs, gt_img) loss_net = (1.0 - opt.lambda_dssim) * Ll1_net + opt.lambda_dssim * (1.0 - ssim(image_gs, gt_img)) # loss_net = (1.0 - opt.lambda_dssim) * Ll1_net loss_net.backward(retain_graph=True) eghuman.opt.step() # parameters_updated = False # for param in eghuman.gs_decoder.parameters(): # if param.grad is not None and param.grad.abs().sum().item() != 0: # parameters_updated = True # break # if parameters_updated: # print("梯度已更新") # else: # print("梯度未更新") # parameters_updated = False # for param in eghuman.tri_encoder.parameters(): # if param.grad is not None and param.grad.abs().sum().item() != 0: # parameters_updated = True # break # if parameters_updated: # print("梯度已更新") # else: # print("梯度未更新") iter_end_net.record() torch.cuda.synchronize() with torch.no_grad(): # Progress bar ema_loss_for_log = 0.4 * loss_net.item() + 0.6 * ema_loss_for_log if net_itr % 1 == 0: training_bar.set_postfix({"Loss": f"{ema_loss_for_log:.{7}f}"}) training_bar.update(1) if net_itr == opt.iterations: training_bar.close() #valid the scene trained ###### # valid_report() training_report(tb_writer, net_itr, Ll1_net, loss_net, l1_loss, iter_start_net.elapsed_time(iter_end), testing_iterations, scene, render, (pipe, background)) if (net_itr in saving_iterations): #save print("\n[ITER {}] Saving Gaussians".format(net_itr)) scene.save(net_itr, 'ply') if dataset.type_3d == 'obj': scene.save(net_itr, dataset.type_3d) if (net_itr in checkpoint_iterations): print("\n[ITER {}] Saving Checkpoint".format(net_itr)) torch.save((gaussians.capture(), net_itr), scene.model_path + "/chkpnt" + str(net_itr) + ".pth") torch.cuda.empty_cache() #save edhuman model print("\n[Person {}] Saving Checkpoint".format(step)) os.makedirs(dataset.eghuamn_model_path, exist_ok=True) torch.save({'e_2d':eghuman.feature_net.state_dict(), 'e_3d':eghuman.tri_encoder.state_dict(), 'd':eghuman.gs_decoder.state_dict()}, dataset.eghuamn_model_path + str(epoch) + '_' + f'{str(step).zfill(5)}.pth') print('person_0 train complete!!!') break epoch = epoch + 1 def prepare_output_and_logger(args): if not args.model_path: if os.getenv('OAR_JOB_ID'): unique_str=os.getenv('OAR_JOB_ID') else: unique_str = str(uuid.uuid4()) args.model_path = os.path.join("./output/", unique_str[0:10]) # Set up output folder print("Output folder: {}".format(args.model_path)) os.makedirs(args.model_path, exist_ok = True) with open(os.path.join(args.model_path, "cfg_args"), 'w') as cfg_log_f: cfg_log_f.write(str(Namespace(**vars(args)))) # Create Tensorboard writer tb_writer = None if TENSORBOARD_FOUND: tb_writer = SummaryWriter(args.model_path) else: print("Tensorboard not available: not logging progress") return tb_writer def training_report(tb_writer, iteration, Ll1, loss, l1_loss, elapsed, testing_iterations, scene : Scene, renderFunc, renderArgs): if tb_writer: tb_writer.add_scalar('train_loss_patches/l1_loss', Ll1.item(), iteration) tb_writer.add_scalar('train_loss_patches/total_loss', loss.item(), iteration) tb_writer.add_scalar('iter_time', elapsed, iteration) # Report test and samples of training set if iteration in testing_iterations: torch.cuda.empty_cache() validation_configs = ({'name': 'test', 'cameras' : scene.getTestCameras()}, {'name': 'train', 'cameras' : [scene.getTrainCameras()[idx % len(scene.getTrainCameras())] for idx in range(5, 30, 5)]}) for config in validation_configs: if config['cameras'] and len(config['cameras']) > 0: l1_test = 0.0 psnr_test = 0.0 for idx, viewpoint in enumerate(config['cameras']): image = torch.clamp(renderFunc(viewpoint, scene.gaussians, *renderArgs)["render"], 0.0, 1.0) gt_image = torch.clamp(viewpoint.original_image.to("cuda:0"), 0.0, 1.0) if tb_writer and (idx < 5): tb_writer.add_images(config['name'] + "_view_{}/render".format(viewpoint.image_name), image[None], global_step=iteration) if iteration == testing_iterations[0]: tb_writer.add_images(config['name'] + "_view_{}/ground_truth".format(viewpoint.image_name), gt_image[None], global_step=iteration) l1_test += l1_loss(image, gt_image).mean().double() psnr_test += psnr(image, gt_image).mean().double() psnr_test /= len(config['cameras']) l1_test /= len(config['cameras']) print("\n[ITER {}] Evaluating {}: L1 {} PSNR {}".format(iteration, config['name'], l1_test, psnr_test)) if tb_writer: tb_writer.add_scalar(config['name'] + '/loss_viewpoint - l1_loss', l1_test, iteration) tb_writer.add_scalar(config['name'] + '/loss_viewpoint - psnr', psnr_test, iteration) if tb_writer: tb_writer.add_histogram("scene/opacity_histogram", scene.gaussians.get_opacity, iteration) tb_writer.add_scalar('total_points', scene.gaussians.get_xyz.shape[0], iteration) torch.cuda.empty_cache() if __name__ == "__main__": # Set up command line argument parser parser = ArgumentParser(description="Training script parameters") lp = ModelParams(parser) op = OptimizationParams(parser) pp = PipelineParams(parser) parser.add_argument('--ip', type=str, default="127.0.0.1") parser.add_argument('--port', type=int, default=6060) parser.add_argument('--debug_from', type=int, default=-1) parser.add_argument('--detect_anomaly', action='store_true', default=False) parser.add_argument("--test_iterations", nargs="+", type=int, default=[3_000, 30_000]) parser.add_argument("--save_iterations", nargs="+", type=int, default=[1, 1_000, 3_000, 5_000, 14_000, 30_000]) parser.add_argument("--quiet", action="store_true") parser.add_argument("--checkpoint_iterations", nargs="+", type=int, default=[0, 500, 1_000, 3_000, 14_000, 30_000]) parser.add_argument("--start_checkpoint", type=str, default = None) args = parser.parse_args(sys.argv[1:]) args.save_iterations.append(args.iterations) print("Optimizing " + args.model_path) # Initialize system state (RNG) safe_state(args.quiet) # Start GUI server, configure and run training network_gui.init(args.ip, args.port) torch.autograd.set_detect_anomaly(args.detect_anomaly) dataset_train = eghuman_dataset(lp.extract(args), train = True) dataloader_train = DataLoader(dataset=dataset_train, batch_size=1, shuffle=False, num_workers=0, collate_fn=collate_fn) training(dataloader_train, lp.extract(args), op.extract(args), pp.extract(args), args.test_iterations, args.save_iterations, args.checkpoint_iterations, args.start_checkpoint, args.debug_from) # All done print("\nTraining complete.")