luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update upscaler training

Browse Source
This commit is contained in:
luni
2023-07-02 01:06:38 +02:00
parent d66e6d5a77
commit 74c72fd9e1
1 changed files with 25 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/upscaler_train.py
+25 -15
View File
@@ -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()