160 lines
5.3 KiB
Python
160 lines
5.3 KiB
Python
import os
|
|
import sys
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
|
|
|
import importlib
|
|
import argparse
|
|
parser = argparse.ArgumentParser(description='')
|
|
parser.add_argument('--name', type=str)
|
|
parser.add_argument('--config', type=str)
|
|
parser.add_argument('--debug', action='store_true')
|
|
args = vars(parser.parse_args())
|
|
|
|
|
|
from configs.iccvw.pyr_64_2 import config
|
|
|
|
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = '5'
|
|
if args['debug']:
|
|
os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
|
|
|
|
import time
|
|
import argparse
|
|
import wandb
|
|
import tqdm
|
|
import random
|
|
import numpy as np
|
|
import plotly.graph_objects as go
|
|
import plotly.express as ex
|
|
from PIL import Image
|
|
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
import torch.nn.functional as F
|
|
from torch.profiler import profile, record_function, ProfilerActivity
|
|
from data.dataset import Dataset
|
|
from utils.checkpoints import load_model
|
|
from utils.utils import InfiniteDataLoader as DataLoader
|
|
from losses.losses import kl_loss
|
|
from custom_utils.custom_utils import UV2Mesh, Accumulator, group_inputs_for_levels
|
|
from smpl.smpl_torch import SMPL
|
|
from models.pyramid_model import Pyramid
|
|
|
|
|
|
|
|
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
|
|
|
|
def train_step(batch):
|
|
model.train()
|
|
# with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
|
|
levels_input, conds = group_inputs_for_levels(batch, config)
|
|
|
|
out_levels = model([levels_input, conds])
|
|
|
|
if config['use_skinning']:
|
|
mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
|
|
else:
|
|
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
|
|
|
|
reconstr_loss_total = 0
|
|
kl_total = 0
|
|
|
|
#for i, res in enumerate(config['pyramid']['resolutions']):
|
|
i = -1
|
|
pred, mean, logvar = out_levels[i]
|
|
|
|
mesh_pred = uv2mesh(pred)
|
|
|
|
if config['use_skinning']:
|
|
shape = batch['pose_shape'][:, 72:]
|
|
gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
|
|
shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
|
|
|
|
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
|
|
mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
|
|
else:
|
|
mesh_pred_posed = mesh_pred
|
|
|
|
|
|
kl = kl_loss(mean, logvar)
|
|
|
|
loss: torch.Tensor = g_balancing_coeff * kl
|
|
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
model.eval()
|
|
with torch.inference_mode():
|
|
metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
|
|
|
|
|
|
return { **{ 'train/loss': loss, 'train/reconstr_loss_total': reconstr_loss_total, 'train/kl_total': kl_total } }
|
|
|
|
|
|
def train():
|
|
global g_balancing_coeff
|
|
global step
|
|
|
|
for batch in train_dataloader:
|
|
train_step(batch)
|
|
|
|
step += 1
|
|
print(step)
|
|
profiler.step()
|
|
|
|
if step > 300:
|
|
break
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
|
|
skirt_modifier = '_skirt' if config['is_skirt'] else ''
|
|
dataset_resolutions = config['pyramid']['resolutions'].copy()
|
|
if 512 not in dataset_resolutions:
|
|
dataset_resolutions.append(512)
|
|
|
|
dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_single.txt', use_cache=True, resolutions=dataset_resolutions, is_cape=config['is_cape'])
|
|
train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True, )
|
|
train_accumulator = Accumulator()
|
|
|
|
model: Pyramid = Pyramid(config)
|
|
|
|
optimizer = optim.Adam(model.parameters(), lr=config['lr'])
|
|
|
|
g_balancing_coeff = config['starting_balancing_coeff']
|
|
reg_coeff = 1
|
|
|
|
step = 0
|
|
if config['pretrained_model'] != '':
|
|
# assert len(config['pyramid']['resolutions']) == 1
|
|
load_model(model, config['pretrained_model'], optimizer=optimizer, iteration=config['pretrained_model_iteration'])
|
|
|
|
if len(config['pyramid']['resolutions']) == 1:
|
|
for param in model.static_enc.parameters(): param.requires_grad = False
|
|
for param in model.body_enc.parameters(): param.requires_grad = False
|
|
for param in model.norm_enc.parameters(): param.requires_grad = False
|
|
else:
|
|
load_model(model, config['pyramid']['pretrained_pyramid'], iteration=config['pyramid']['pretrained_pyramid_iteration'])
|
|
for param in model.static_enc.parameters(): param.requires_grad = False
|
|
for param in model.body_enc.parameters(): param.requires_grad = False
|
|
for param in model.norm_enc.parameters(): param.requires_grad = False
|
|
|
|
for i in range(len(model.levels)-1):
|
|
for param in model.levels[i].parameters(): param.requires_grad = False
|
|
|
|
# create a profiler object
|
|
profiler = torch.profiler.profile(
|
|
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
|
|
schedule=torch.profiler.schedule(wait=2, warmup=5, active=3, repeat=2),
|
|
on_trace_ready=torch.profiler.tensorboard_trace_handler("log_dir"),
|
|
)
|
|
|
|
with profiler:
|
|
train()
|
|
|
|
profiler.stop()
|
|
print(profiler.key_averages().table(sort_by="cpu_time_total", row_limit=25)) |