From 74c72fd9e1fcf5921aa82751bfabd4ed23438886 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sun, 2 Jul 2023 01:06:38 +0200 Subject: [PATCH] update upscaler training --- scripts/upscaler_train.py | 40 ++++++++++++++++++++++++--------------- 1 file changed, 25 insertions(+), 15 deletions(-) diff --git a/scripts/upscaler_train.py b/scripts/upscaler_train.py index 7199a1d..f581ecd 100644 --- a/scripts/upscaler_train.py +++ b/scripts/upscaler_train.py @@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true') args = vars(parser.parse_args()) os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" -os.environ["CUDA_VISIBLE_DEVICES"] = '6' +os.environ["CUDA_VISIBLE_DEVICES"] = '2' import time @@ -49,26 +49,36 @@ gender_modifier = '' if not is_cape else ('_m' if is_male else '_f') smpl_res = 3638 if is_cape else 14475 use_cache = False name=f'iccv_{low_res}' + skirt_modifier + cape_modifier + gender_modifier +train_dir = '/data-local/127/cloth3d_processed' +valid_dir = '/data-local/127/cloth3d_processed' -if not is_skirt: - train_dir = '/data-local/127/train' - valid_dir = '/data-local/127/valid' +if is_skirt: + train_frames = 'frames_train_skirt.txt' + valid_frames = 'frames_valid_skirt.txt' else: - train_dir = '/data-local/128/train_skirt' - valid_dir = '/data-local/127/valid_skirt' + train_frames = 'frames_train.txt' + valid_frames = 'frames_valid.txt' + +if is_cape: + train_frames = None + valid_frames = None if is_cape: if is_male: - train_dir = '/data-local/127/cape_processed/train/male' - valid_dir = '/data-local/127/cape_processed/test/male' + train_dir = '/data-local/127_2/cape_processed/train/male' + valid_dir = '/data-local/127_2/cape_processed/test/male' else: - train_dir = '/data-local/127/cape_processed/train/female' - valid_dir = '/data-local/127/cape_processed/test/female' + train_dir = '/data-local/127_2/cape_processed/train/female' + valid_dir = '/data-local/127_2/cape_processed/test/female' +print(f'nohup python scripts/upscaler_train.py --name {name} > nohup/up/up_{name}.out 2>&1 &') +if args['debug']: + exit() def init_wandb(): - run = wandb.init(project='upscale', name=name, notes='', dir='/data-local/127/temp', tags=['upscaler', 'iccv']) + wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True) + run = wandb.init(project='upscale', name=name, notes='', dir='/data-local/127/temp', tags=['upscaler', 'iccvw']) artifact = wandb.Artifact('code', type='code') artifact.add_file('scripts/upscaler_train.py') artifact.add_file('losses/losses.py') @@ -255,8 +265,8 @@ if __name__ == '__main__': } # dataset = Dataset('/data-local/127/train', config['signature'], resolutions=[512], device=device, frames_file='frames_jumpsuit.txt') - dataset_train = Dataset(train_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=None, resolutions=[512, low_res], is_cape=is_cape) - dataset_valid = Dataset(valid_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=None, resolutions=[512, low_res], is_cape=is_cape) + dataset_train = Dataset(train_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=train_frames, resolutions=[512, low_res], is_cape=is_cape) + dataset_valid = Dataset(valid_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=valid_frames, resolutions=[512, low_res], is_cape=is_cape) train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True) valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True) train_accumulator = Accumulator() @@ -279,13 +289,13 @@ if __name__ == '__main__': _MSSSIM_small = MS_SSIM(data_range=2, size_average=True, channel=3, weights=[1.0]) step = 0 - ckpt_saver = CheckpointSaver('checkpoints/iccv/torch/' + name, top_n=5) + ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + name, top_n=5) print(f'started training: {name}') if not args['debug']: # model = torch.compile(model) - wandb.watch(model, log='all', log_freq=1000, log_graph=True) + wandb.watch(model, log='all', log_freq=10000, log_graph=True) train()