started training pyramid

This commit is contained in:
2023-02-26 22:25:50 +00:00
parent a678724711
commit 64ecb900f5
8 changed files with 314 additions and 79 deletions
+16 -7
View File
@@ -36,14 +36,17 @@ from models.upscaler import UpscalerBasicConv as Upscaler
# !!!!!!!!!! CHANGE SKIRT OR NOT
low_res = 128
low_res = 256
is_skirt = True
is_cape = False
is_male = False
skirt_modifier = '_skirt' if is_skirt else ''
cape_modifier = '_cape' if is_cape else ''
gender_modifier = '' if not is_cape else ('_m' if is_male else '_f')
smpl_res = 3638 if is_cape else 14475
name=f'iccv_{low_res}_2' + skirt_modifier
name=f'iccv_{low_res}_2' + skirt_modifier + cape_modifier + gender_modifier
def init_wandb():
run = wandb.init(project='upscale', name=name,
notes='\
@@ -202,7 +205,6 @@ def train():
if step % 5000 == 0:
manager.save(checkpoint_number=step)
wandb.log(to_log, step = step)
print(step)
@@ -255,10 +257,17 @@ if __name__ == '__main__':
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/valid_skirt', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
if is_cape:
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/train/male', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/cape_processed/train/male', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/test/male', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/cape_processed/test/male', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
if is_male:
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/train/male', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/cape_processed/train/male', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/test/male', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/cape_processed/test/male', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
else:
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/train/female', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/cape_processed/train/female', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/test/female', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/cape_processed/test/female', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
if low_res == 32 or low_res == 64:
model = Upscaler(low_res=low_res, kernel_size=3)