From 32c3d1103fb7716388e24f3a99346d2405517f35 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 2 Jun 2023 09:36:01 +0000 Subject: [PATCH] train with fixed reg losses --- configs/direct_7.py | 12 ++++++------ scripts/direct_train.py | 24 +++++++++++++----------- 2 files changed, 19 insertions(+), 17 deletions(-) diff --git a/configs/direct_7.py b/configs/direct_7.py index d0f2404..37dade0 100644 --- a/configs/direct_7.py +++ b/configs/direct_7.py @@ -2,9 +2,9 @@ from utils.convnext_utils import CNBlockConfig config = dict() -config['gpu'] = '6' -config['name'] = 'direct_7' -config['notes'] = 'with normals input but only single garment sequence with msssim and l2 mesh loss only and without skinning' +config['gpu'] = '0' +config['name'] = 'direct_9' +config['notes'] = 'with normals input whole dataset with msssim and l2 mesh loss and reg losses and without skinning' config['tags'] = ['torch', 'direct'] config['script'] = 'scripts/direct_train.py' @@ -20,12 +20,12 @@ config['batch_size'] = 4 config['train_dir'] = '/data-local/127/train' config['valid_dir'] = '/data-local/127/train' config['test_dir'] = '/data-net/dataset/test_ready' -config['train_frames'] = 'frames_00761.txt' -config['valid_frames'] = 'frames_00761.txt' +config['train_frames'] = None +config['valid_frames'] = None config['resolution'] = 512 config['is_skirt'] = False -config['dataset_in_memory'] = True +config['dataset_in_memory'] = False config['is_cape'] = False config['remove_global_orientation'] = True config['device'] = 'cuda' diff --git a/scripts/direct_train.py b/scripts/direct_train.py index 574912b..5b18b46 100644 --- a/scripts/direct_train.py +++ b/scripts/direct_train.py @@ -13,7 +13,7 @@ args = vars(parser.parse_args()) config_module = importlib.import_module(args['config']) config = config_module.config if args['debug']: - from configs.direct_8 import config + from configs.direct_7 import config # from configs.ae_static_orig_2 import config os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] @@ -98,8 +98,15 @@ def train_step(batch): l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1)) + e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) + c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) + c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio) + # l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2])) + n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) + + ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1)) - loss = l2_mesh_loss + (1 - ms_ssim) * 1e3 + loss = l2_mesh_loss + (1 - ms_ssim) * 1e3 + (e_loss + c_loss + n_loss) * config['reg_loss_coeff'] optimizer.zero_grad() loss.backward() @@ -112,12 +119,7 @@ def train_step(batch): 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']) - e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) - c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) - c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio) - # l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2])) - n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) - + return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/euclid_error': metric, 'train/collision_ratio': c_ratio, 'train/ms_ssim': ms_ssim, 'train/collision_loss': c_loss, 'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss } @@ -157,7 +159,7 @@ def valid_step(batch): n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1)) - loss = l2_mesh_loss + (1 - ms_ssim) * 1e3 + loss = l2_mesh_loss + (1 - ms_ssim) * 1e3 + (e_loss + c_loss + n_loss) * config['reg_loss_coeff'] metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_mask']) @@ -225,8 +227,8 @@ def plot_everything(batch: dict, type: str): collision_fixed = [np.concatenate((body[i], garment_posed_fixed[i]), axis=0) for i in range(len(garment_posed))] collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(garment_posed))] - faces_ = np.load('/data-net/dataset/00761/faces/00761.npy') - faces = [faces_ for i in range(len(garment_posed))] + # faces_ = np.load('/data-net/dataset/00761/faces/00761.npy') + # faces = [faces_ for i in range(len(garment_posed))] collision_faces = [np.concatenate((F, faces[i] + config['smpl_res']), axis=0) for i in range(len(garment_posed))] figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }