small updates
This commit is contained in:
+20
-62
@@ -10,10 +10,12 @@ 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.cvae_norm_1 import config
|
||||
from configs.iccvw.cvae_norm_1 import config
|
||||
else:
|
||||
config_module = importlib.import_module(args['config'])
|
||||
config = config_module.config
|
||||
|
||||
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
|
||||
@@ -51,7 +53,8 @@ from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary
|
||||
|
||||
|
||||
def init_wandb():
|
||||
run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
|
||||
wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
|
||||
run = wandb.init(project='iccvw', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
|
||||
artifact = wandb.Artifact('code', type='code')
|
||||
artifact.add_file('scripts/cvae_norm.py')
|
||||
artifact.add_file('losses/losses.py')
|
||||
@@ -80,13 +83,13 @@ def kl_loss(mean: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
|
||||
def train_step(batch):
|
||||
model.train()
|
||||
pred, mean, logvar = model([batch['uv_normals'], batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']])
|
||||
pred = pred * 1.1
|
||||
pred = pred * 1.4
|
||||
mesh_pred = uv2mesh(pred)
|
||||
|
||||
|
||||
l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_normals'] - pred) ** 2), axis=-1), axis=[1, 2]))
|
||||
l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_normals'] - mesh_pred) ** 2, axis=-1), 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))
|
||||
kl = kl_loss(mean, logvar)
|
||||
|
||||
reconstr_loss = l2_loss + l2_mesh_loss + (1 - ms_ssim)
|
||||
@@ -114,12 +117,12 @@ def valid_step(batch):
|
||||
model.eval()
|
||||
with torch.inference_mode():
|
||||
pred, mean, logvar = model([batch['uv_normals'], batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']])
|
||||
pred = pred * 1.1
|
||||
pred = pred * 1.4
|
||||
mesh_pred = uv2mesh(pred)
|
||||
|
||||
l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_normals'] - pred) ** 2), axis=-1), axis=[1, 2]))
|
||||
l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_pred - batch['mesh_normals']) ** 2, axis=-1), axis=-1))
|
||||
ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_static'] + 1.1))
|
||||
ms_ssim = torch.mean(_MSSSIM(pred + 1.4, batch['uv_static'] + 1.4))
|
||||
kl = kl_loss(mean, logvar)
|
||||
|
||||
reconstr_loss = l2_loss + l2_mesh_loss + (1 - ms_ssim)
|
||||
@@ -142,12 +145,12 @@ def plot_everything(batch: dict, type: str):
|
||||
with torch.inference_mode():
|
||||
pred, mean, logvar = model([batch['uv_normals'], batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']])
|
||||
|
||||
pred = pred * 1.1
|
||||
pred = pred * 1.4
|
||||
# mesh_pred = uv2mesh(pred).detach().cpu().numpy()
|
||||
|
||||
uv_static_img = torch.permute(batch['uv_static'], [0,2,3,1])
|
||||
uv_body_posed_img = torch.permute(batch['uv_body_posed'], [0,2,3,1])
|
||||
orig_img = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.1) / 2.2
|
||||
orig_img = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.4) / 2.8
|
||||
pred_img = torch.permute(pred, [0,2,3,1])
|
||||
img = torch.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
|
||||
imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)]
|
||||
@@ -281,54 +284,6 @@ def sample(folder):
|
||||
|
||||
|
||||
|
||||
def test():
|
||||
str_split = config['name'].split('_')
|
||||
test_idx = str_split.index('test')
|
||||
checkpoint_name = '_'.join(str_split[:test_idx])
|
||||
load_model(model, 'checkpoints/' + checkpoint_name)
|
||||
|
||||
dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False, do_repeat=False)
|
||||
static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False)
|
||||
|
||||
step = 0
|
||||
losses = []
|
||||
|
||||
for batch in tqdm.tqdm(dataset_test):
|
||||
batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
|
||||
|
||||
test_loss = test_step(batch)
|
||||
losses.append(test_loss)
|
||||
|
||||
to_log = { **{ 'current_step': step }, **test_loss }
|
||||
if step % 1000 == 0:
|
||||
figs = plot_everything(batch, 'test')
|
||||
to_log = { **to_log, **figs }
|
||||
|
||||
wandb.log(to_log, step = step)
|
||||
|
||||
# print(step)
|
||||
step += 1
|
||||
|
||||
# averages = dict()
|
||||
path_to_plotly_html = "temp/plotly_figure.html"
|
||||
|
||||
histograms = dict()
|
||||
for key in test_loss.keys():
|
||||
accumulated = [losses[i][key].numpy() for i in range(len(losses))]
|
||||
fig = ex.histogram(accumulated, title='Average: ' + str(np.asarray(accumulated).mean()))
|
||||
fig.write_html(path_to_plotly_html, auto_play = False)
|
||||
|
||||
histograms['all/' + key.split('/')[-1] + '_hist'] = wandb.Html(path_to_plotly_html)
|
||||
# averages['all/' + key.split('/')[-1] + '_avg'] = np.asarray(accumulated).mean()
|
||||
|
||||
|
||||
|
||||
wandb.log({ **histograms }, step = step)
|
||||
wandb.finish()
|
||||
pass
|
||||
|
||||
|
||||
|
||||
def train():
|
||||
global g_balancing_coeff
|
||||
global step
|
||||
@@ -358,10 +313,10 @@ def train():
|
||||
valid_accumulator.reset()
|
||||
print(step)
|
||||
|
||||
if train_loss['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
|
||||
if train_loss['train/euclid_error'] < config['balancing_activation_limit']:
|
||||
g_balancing_coeff += 1 / 10000
|
||||
|
||||
if step % 10000 == 0:
|
||||
if step % 20000 == 0:
|
||||
ckpt_saver(model, step, optimizer, None)
|
||||
|
||||
step += 1
|
||||
@@ -387,11 +342,14 @@ if __name__ == '__main__':
|
||||
_MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
|
||||
g_balancing_coeff = config['starting_balancing_coeff']
|
||||
|
||||
ckpt_saver = CheckpointSaver('checkpoints/iccv/' + config['name'], top_n=5)
|
||||
ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=25)
|
||||
|
||||
step = 0
|
||||
if config['pretrained_model'] != '':
|
||||
step = load_model(model, config['pretrained_model'])
|
||||
if 'ignore_saved_optimizer' in config.keys() and config['ignore_saved_optimizer']:
|
||||
step = load_model(model, config['pretrained_model'])
|
||||
else:
|
||||
step = load_model(model, config['pretrained_model'], optimizer=optimizer)
|
||||
|
||||
if not args['debug']:
|
||||
# model = torch.compile(model)
|
||||
|
||||
Reference in New Issue
Block a user