gha / data_tool /train_generalization.py

Add files using upload-large-folder tool

3217786 verified 3 months ago

21.9 kB

#
# 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.")

	#
	# 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*inp.pi*xyz_sample))
	posison_encoding.append(torch.cos(2*inp.pi*xyz_sample))
	posison_encoding.append(torch.sin(2*Lnp.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.")