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

ae_static for cape

Browse Source
This commit is contained in:
luni
2023-05-21 22:15:36 +02:00
parent 70162cf8ac
commit 28d1ad6469
2 changed files with 106 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
configs/ae_static_cape_m.py
+74
View File
@@ -0,0 +1,74 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '5'
config['name'] = 'static_torch_cape_m'
config['notes'] = ' '
config['tags'] = ['torch', 'static']
config['script'] = 'scripts/ae_static_train.py'
config['pretrained_model'] = ''
config['batch_size'] = 4
config['train_dir'] = '/data-local/127/cape_processed/train/male'
config['valid_dir'] = '/data-local/127/cape_processed/test/male'
config['test_dir'] = '/data-local/127/cape_processed/test/male'
config['train_frames'] = None
config['valid_frames'] = None
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = True
config['is_cape'] = True
config['remove_global_orientation'] = True
config['device'] = 'cuda'
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['latent_dim'] = 128
config['bottleneck_scale'] = 2
config['cond_ae'] = dict()
config['cond_ae']['device'] = config['device']
config['cond_ae']['is_label'] = False
config['cond_ae']['latent_dim'] = config['latent_dim']
config['cond_ae']['kernel_size'] = 7
config['cond_ae']['bottleneck_size'] = 8 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(8, 16, 3),
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, config['cond_ae']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 8, 3),
CNBlockConfig(8, 3, 3),
# CNBlockConfig(3, 3, 3),
]
config['cond_ae']['encoder_block_settings'] = block_setting
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
config['cond_ae']['activation'] = 'tanh'
config['lr'] = 0.0005
config['cond_ae']['upscale_output'] = False
config['pretrained_model'] = ''
config['signature'] = { 'id',
'frame_nr',
'uv_static',
'mesh_static',
'mesh_mask',
'uv_mask',
}
scripts/ae_static_train.py
+32 -19
View File
@@ -7,14 +7,18 @@ import argparse
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
if args['debug']:
from configs.ae_static_cape_m 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']
if args['debug']:
os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
import time
import argparse
@@ -66,6 +70,7 @@ uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_
def train_step(batch):
model.train()
# with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
pred = model([batch['uv_static'], batch['uv_mask']])
ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1))
@@ -75,6 +80,10 @@ def train_step(batch):
optimizer.zero_grad()
loss.backward()
optimizer.step()
# scaler.scale(loss).backward()
# scaler.step(optimizer)
# scaler.update()
# optimizer.zero_grad()
model.eval()
with torch.inference_mode():
@@ -113,7 +122,10 @@ def plot_everything(batch: dict, type: str):
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_static']
mesh_pred = torch.cat((mesh_pred[:, :, 0:1], -mesh_pred[:, :, 2:3], mesh_pred[:, :, 1:2]), dim=2)
mesh_orig = torch.cat((mesh_orig[:, :, 0:1], -mesh_orig[:, :, 2:3], mesh_orig[:, :, 1:2]), dim=2)
orig_img = torch.permute(batch['uv_static'], [0,2,3,1])
pred_img = torch.permute(pred, [0,2,3,1])
img = torch.concat((orig_img, pred_img), axis=2)
@@ -157,34 +169,29 @@ 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, scaler)
ckpt_saver_dec(model.decoder, step, optimizer, scaler)
step += 1
if __name__ == '__main__':
init_wandb()
# device = "cuda" if torch.cuda.is_available() else "cpu"
# parser = argparse.ArgumentParser(description='')
# parser.add_argument('--name', type=str)
# args = vars(parser.parse_args())
if not args['debug']:
init_wandb()
resolution_modifier = ''
if config['resolution'] != 512:
resolution_modifier = '_' + str(config['resolution'])
use_amp = True
skirt_modifier = '_skirt' if config['is_skirt'] else ''
# dataset = Dataset('/data-local/127/train', config['signature'], resolutions=[512], device=device, frames_file='frames_jumpsuit.txt')
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()
@@ -192,6 +199,7 @@ if __name__ == '__main__':
model : AutoEncoder = AutoEncoder(config['cond_ae'])
optimizer = optim.Adam(model.parameters(), lr=config['lr'])
scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
# _MSSSIM = MSSSIM(data_range=2).to(config['device'])
_MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
@@ -200,11 +208,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()