Spaces:
Sleeping
Sleeping
File size: 5,003 Bytes
ec9a6bc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
import os
import torch
import argparse
from config.config import config_train
from lib.dataset.Dataset import GaussianDataset
from lib.dataset.DataLoaderX import DataLoaderX
from lib.module.MeshHeadModule import MeshHeadModule
from lib.module.GaussianHeadModule import GaussianHeadModule
from lib.module.SuperResolutionModule import SuperResolutionModule
from lib.module.CameraModule import CameraModule
from lib.recorder.Recorder import GaussianHeadTrainRecorder
from lib.trainer.GaussianHeadTrainer import GaussianHeadTrainer
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--config', type=str, default='config/train_s2_N031.yaml')
arg = parser.parse_args()
cfg = config_train()
cfg.load(arg.config)
cfg = cfg.get_cfg()
dataset = GaussianDataset(cfg.dataset)
dataloader = DataLoaderX(dataset, batch_size=cfg.batch_size, shuffle=True, pin_memory=True, drop_last=True, num_workers=8)
device = torch.device('cuda:%d' % cfg.gpu_id)
torch.cuda.set_device(cfg.gpu_id)
if os.path.exists(cfg.load_gaussianhead_checkpoint):
gaussianhead_state_dict = torch.load(cfg.load_gaussianhead_checkpoint, map_location=lambda storage, loc: storage)
gaussianhead = GaussianHeadModule(cfg.gaussianheadmodule,
xyz=gaussianhead_state_dict['xyz'],
feature=gaussianhead_state_dict['feature'],
landmarks_3d_neutral=gaussianhead_state_dict['landmarks_3d_neutral']).to(device)
gaussianhead.load_state_dict(gaussianhead_state_dict)
else:
meshhead_state_dict = torch.load(cfg.load_meshhead_checkpoint, map_location=lambda storage, loc: storage)
meshhead = MeshHeadModule(cfg.meshheadmodule, meshhead_state_dict['landmarks_3d_neutral']).to(device)
meshhead.load_state_dict(meshhead_state_dict)
meshhead.subdivide()
with torch.no_grad():
data = meshhead.reconstruct_neutral()
gaussianhead = GaussianHeadModule(cfg.gaussianheadmodule,
xyz=data['verts'].cpu(),
feature=torch.atanh(data['verts_feature'].cpu()),
landmarks_3d_neutral=meshhead.landmarks_3d_neutral.detach().cpu(),
add_mouth_points=True).to(device)
gaussianhead.exp_color_mlp.load_state_dict(meshhead.exp_color_mlp.state_dict())
gaussianhead.pose_color_mlp.load_state_dict(meshhead.pose_color_mlp.state_dict())
gaussianhead.exp_deform_mlp.load_state_dict(meshhead.exp_deform_mlp.state_dict())
gaussianhead.pose_deform_mlp.load_state_dict(meshhead.pose_deform_mlp.state_dict())
supres = SuperResolutionModule(cfg.supresmodule).to(device)
if os.path.exists(cfg.load_supres_checkpoint):
supres.load_state_dict(torch.load(cfg.load_supres_checkpoint, map_location=lambda storage, loc: storage))
camera = CameraModule()
recorder = GaussianHeadTrainRecorder(cfg.recorder)
optimized_parameters = [{'params' : supres.parameters(), 'lr' : cfg.lr_net},
{'params' : gaussianhead.xyz, 'lr' : cfg.lr_net * 0.1},
{'params' : gaussianhead.feature, 'lr' : cfg.lr_net * 0.1},
{'params' : gaussianhead.exp_color_mlp.parameters(), 'lr' : cfg.lr_net},
{'params' : gaussianhead.pose_color_mlp.parameters(), 'lr' : cfg.lr_net},
{'params' : gaussianhead.exp_deform_mlp.parameters(), 'lr' : cfg.lr_net},
{'params' : gaussianhead.pose_deform_mlp.parameters(), 'lr' : cfg.lr_net},
{'params' : gaussianhead.exp_attributes_mlp.parameters(), 'lr' : cfg.lr_net},
{'params' : gaussianhead.pose_attributes_mlp.parameters(), 'lr' : cfg.lr_net},
{'params' : gaussianhead.scales, 'lr' : cfg.lr_net * 0.3},
{'params' : gaussianhead.rotation, 'lr' : cfg.lr_net * 0.1},
{'params' : gaussianhead.opacity, 'lr' : cfg.lr_net}]
if os.path.exists(cfg.load_delta_poses_checkpoint):
delta_poses = torch.load(cfg.load_delta_poses_checkpoint)
else:
delta_poses = torch.zeros([dataset.num_exp_id, 6]).to(device)
if cfg.optimize_pose:
delta_poses = delta_poses.requires_grad_(True)
optimized_parameters.append({'params' : delta_poses, 'lr' : cfg.lr_pose})
else:
delta_poses = delta_poses.requires_grad_(False)
optimizer = torch.optim.Adam(optimized_parameters)
trainer = GaussianHeadTrainer(dataloader, delta_poses, gaussianhead, supres, camera, optimizer, recorder, cfg.gpu_id)
trainer.train(0, 1000)
|