update training scripts
This commit is contained in:
@@ -14,7 +14,7 @@ args = vars(parser.parse_args())
|
||||
config_module = importlib.import_module(args['config'])
|
||||
config = config_module.config
|
||||
if args['debug']:
|
||||
from configs.ae_body_cape_m import config
|
||||
from configs.iccvw.ae_body_cape_m import config
|
||||
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
|
||||
if args['debug']:
|
||||
@@ -49,6 +49,7 @@ from pytorch_msssim import MS_SSIM
|
||||
|
||||
|
||||
def init_wandb():
|
||||
wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
|
||||
run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
|
||||
artifact = wandb.Artifact('code', type='code')
|
||||
artifact.add_file('scripts/ae_body_train.py')
|
||||
@@ -198,15 +199,15 @@ 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/torch/' + config['name'] + '_enc', top_n=5)
|
||||
ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_dec', top_n=5)
|
||||
ckpt_saver_enc = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_enc', top_n=5)
|
||||
ckpt_saver_dec = CheckpointSaver('checkpoints/iccvw/' + 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)
|
||||
wandb.watch(model, log='all', log_freq=10000, log_graph=True)
|
||||
|
||||
train()
|
||||
|
||||
|
||||
@@ -49,7 +49,8 @@ from pytorch_msssim import MS_SSIM
|
||||
|
||||
|
||||
def init_wandb():
|
||||
run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127_2/temp', tags=config['tags'])
|
||||
wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
|
||||
run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
|
||||
artifact = wandb.Artifact('code', type='code')
|
||||
artifact.add_file('scripts/ae_norm_train.py')
|
||||
artifact.add_file('losses/losses.py')
|
||||
@@ -70,9 +71,9 @@ uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_
|
||||
|
||||
def train_step(batch):
|
||||
model.train()
|
||||
pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.1
|
||||
pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.4
|
||||
|
||||
ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_normals'] + 1.1))
|
||||
ms_ssim = torch.mean(_MSSSIM(pred + 1.4, batch['uv_normals'] + 1.4))
|
||||
loss = 1e3 * (1 - ms_ssim)
|
||||
|
||||
optimizer.zero_grad()
|
||||
@@ -93,13 +94,13 @@ def valid_step(batch):
|
||||
model.eval()
|
||||
|
||||
with torch.inference_mode():
|
||||
pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.1
|
||||
pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.4
|
||||
|
||||
mesh_pred = uv2mesh(pred)
|
||||
mesh_orig = batch['mesh_normals']
|
||||
|
||||
l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_normals'] - pred)), axis=1))
|
||||
ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_normals'] + 1.1))
|
||||
ms_ssim = torch.mean(_MSSSIM(pred + 1.4, batch['uv_normals'] + 1.4))
|
||||
loss = 1e3 * (1 - ms_ssim)
|
||||
|
||||
metric = torch.mean(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1)))
|
||||
@@ -109,7 +110,7 @@ def valid_step(batch):
|
||||
|
||||
def plot_everything(batch: dict, type: str):
|
||||
model.eval()
|
||||
pred = model([batch['uv_normals'], batch['uv_mask']]).detach() * 1.1
|
||||
pred = model([batch['uv_normals'], batch['uv_mask']]).detach() * 1.4
|
||||
|
||||
# mesh_pred = uv2mesh(pred)
|
||||
# mesh_orig = batch['mesh_normals']
|
||||
@@ -117,7 +118,7 @@ def plot_everything(batch: dict, type: str):
|
||||
orig_img = torch.permute(batch['uv_normals'], [0,2,3,1])
|
||||
pred_img = torch.permute(pred, [0,2,3,1])
|
||||
img = torch.concat((orig_img, pred_img), axis=2)
|
||||
imgs = [wandb.Image((img[i].cpu().numpy() + 1.1) / 2.2) for i in range(4)]
|
||||
imgs = [wandb.Image((img[i].cpu().numpy() + 1.4) / 2.8) for i in range(4)]
|
||||
|
||||
# F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
|
||||
# garment_unposed, faces = mesh_pred.cpu().numpy(), F
|
||||
@@ -196,7 +197,7 @@ if __name__ == '__main__':
|
||||
optimizer = optim.Adam(model.parameters(), lr=config['lr'])
|
||||
|
||||
# _MSSSIM = MSSSIM(data_range=2).to(config['device'])
|
||||
_MSSSIM = MS_SSIM(data_range=2.2, size_average=True, channel=3)
|
||||
_MSSSIM = MS_SSIM(data_range=2.8, size_average=True, channel=3)
|
||||
|
||||
if config['pretrained_model'] != '':
|
||||
load_model(model.encoder, config['pretrained_model'] + '_enc')
|
||||
|
||||
@@ -13,8 +13,8 @@ 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
|
||||
# from configs.ae_static_cape_m import config
|
||||
from configs.iccvw.ae_static_1 import config
|
||||
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
|
||||
if args['debug']:
|
||||
@@ -49,6 +49,7 @@ from pytorch_msssim import MS_SSIM
|
||||
|
||||
|
||||
def init_wandb():
|
||||
wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
|
||||
run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
|
||||
artifact = wandb.Artifact('code', type='code')
|
||||
artifact.add_file('scripts/ae_static_train.py')
|
||||
@@ -208,15 +209,15 @@ 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/torch/' + config['name'] + '_enc', top_n=5)
|
||||
ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_dec', top_n=5)
|
||||
ckpt_saver_enc = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_enc', top_n=5)
|
||||
ckpt_saver_dec = CheckpointSaver('checkpoints/iccvw/' + 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)
|
||||
wandb.watch(model, log='all', log_freq=10000, log_graph=True)
|
||||
|
||||
train()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user