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

ae_curv for cape

Browse Source
This commit is contained in:
luni
2023-05-21 22:37:08 +02:00
parent 28d1ad6469
commit 8e0d061104
2 changed files with 92 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/ae_curv_train.py
+20 -18
View File
@@ -8,14 +8,17 @@ import importlib
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())
config_module = importlib.import_module(args['config'])
config = config_module.config
# from configs.ae_curv_1 import config
if args['debug']:
from configs.ae_curv_cape_m import config
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
if args['debug']:
os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
import time
import wandb
@@ -160,37 +163,31 @@ def train():
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
wandb.log(to_log, step=step)
if not args['debug']:
wandb.log(to_log, step=step)
train_accumulator.reset()
valid_accumulator.reset()
print(step)
if step % 10000 == 0:
ckpt_saver_enc(model.encoder, step, optimizer)
ckpt_saver_dec(model.decoder, step, optimizer)
ckpt_saver_enc(model.encoder, step, optimizer, None)
ckpt_saver_dec(model.decoder, step, optimizer, None)
step += 1
if __name__ == '__main__':
init_wandb()
# device = "cuda" if torch.cuda.is_available() else "cpu"
if not args['debug']:
init_wandb()
# parser = argparse.ArgumentParser(description='')
# parser.add_argument('--name', type=str)
# args = vars(parser.parse_args())
resolution_modifier = ''
if config['resolution'] != 512:
resolution_modifier = '_' + str(config['resolution'])
skirt_modifier = '_skirt' if config['is_skirt'] else ''
uv_mask_body = np.load('custom_utils/mask_256x512' + ('_skirt' if config['is_skirt'] else '') + '.npy')
uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
uv_mask_body = torch.from_numpy(uv_mask_body).float().to(config['device'])
dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'])
dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'])
dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'], is_cape=config['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()
@@ -210,11 +207,16 @@ if __name__ == '__main__':
load_model(model.encoder, config['pretrained_model'] + '_enc')
load_model(model.decoder, config['pretrained_model'] + '_dec')
ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/' + config['name'] + '_enc', top_n=5)
ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/' + config['name'] + '_dec', top_n=5)
ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_enc', top_n=5)
ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_dec', top_n=5)
run_name = config['name']
print(f'started training: {run_name}')
if not args['debug']:
# model = torch.compile(model)
wandb.watch(model, log='all', log_freq=1000, log_graph=True)
train()