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

pyramid training working

Browse Source
This commit is contained in:
luni
2023-07-02 23:23:41 +02:00
parent 2fc0177d5a
commit 592b8df0f0
4 changed files with 668 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/cvae_pyramid.py
+50 -16
View File
@@ -10,11 +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.pyr_32 import config
# from configs.pyr_64_cape_m import config
from configs.iccvw.pyr_32 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']
if args['debug']:
@@ -42,6 +43,7 @@ import torch.nn.functional as F
from data.dataset import Dataset
from utils.checkpoints import CheckpointSaver, load_model
from utils.utils import InfiniteDataLoader as DataLoader
from utils.utils import text_to_html
from utils.postprocess import fix_collision
from losses.losses import collision_loss, normal_loss, kl_loss
from pytorch_msssim import MS_SSIM
@@ -52,7 +54,8 @@ from models.pyramid_model import Pyramid
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_pyramid.py')
artifact.add_file('losses/losses.py')
@@ -99,8 +102,8 @@ def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, col
fig.add_trace(go.Mesh3d(
x=mesh_posed[:,0],
y=mesh_posed[:,1],
z=mesh_posed[:,2],
y=mesh_posed[:,2],
z=mesh_posed[:,1],
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
@@ -270,13 +273,37 @@ def train_step(batch):
reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff']
kl_total += kl
loss = reconstr_loss_total + g_balancing_coeff * kl_total
loss: torch.Tensor = reconstr_loss_total + g_balancing_coeff * kl_total
# scaler.scale(loss).backward()
# scaler.step(optimizer)
# scaler.update()
# optimizer.zero_grad()
# b = 0
# from custom_utils.custom_utils import _maskFaces
# mesh_mask = batch['mesh_mask'][b].cpu().numpy()
# faces = np.load('smpl/smpl_faces_skirt.npy') if batch['garment'][b] in ['Skirt', 'Dress'] else np.load('smpl/smpl_faces.npy')
# faces_body = np.load('smpl/smpl_faces.npy')
# _, faces = _maskFaces(np.zeros((14475, 3)), faces, mesh_mask)
# # fig = go.Figure(data=[go.Mesh3d(x=batch['mesh_posed'].cpu().numpy()[0,mesh_mask,0], y=batch['mesh_posed'].cpu().numpy()[0,mesh_mask,2], z=batch['mesh_posed'].cpu().numpy()[0,mesh_mask,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
# # go.Mesh3d(x=batch['mesh_body_posed'].cpu().numpy()[0,:,0], y=batch['mesh_body_posed'].cpu().numpy()[0,:,2], z=batch['mesh_body_posed'].cpu().numpy()[0,:,1], i=faces_body[:,0], j=faces_body[:,1], k=faces_body[:,2], color='lightblue', opacity=1),
# # ])
# fig = go.Figure(data=[go.Mesh3d(x=mesh_orig.cpu().numpy()[0,mesh_mask,0], y=mesh_orig.cpu().numpy()[0,mesh_mask,2], z=mesh_orig.cpu().numpy()[0,mesh_mask,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
# ])
# fig.update_layout(scene_aspectmode='cube')
# fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
# fig.show()
# ex.imshow((uv_orig[b].cpu().numpy().transpose(1,2,0) + 1) / 2).show()
# upscaled = model.levels[0].enhance(levels_input[0][0], levels_input[0][-2]) # levels_input[i][-2]=uv_mask_res
# upscaled = model.resizers[512](upscaled)
# ex.imshow((levels_input[0][0][b].cpu().numpy().transpose(1,2,0) + 1) / 2).show()
# ex.imshow((upscaled[b].cpu().numpy().transpose(1,2,0) + 1) / 2).show()
optimizer.zero_grad()
loss.backward()
optimizer.step()
@@ -452,12 +479,12 @@ def plot_everything(batch: dict, type: str):
orig_img = (torch.permute(uv_orig, [0,2,3,1]) + 1) / 2
uv_static_img = (torch.permute(batch['uv_static'], [0,2,3,1]) + 1) / 2
uv_body_posed_img = (torch.permute(batch['uv_body_posed'], [0,2,3,1]) + 1) / 2
uv_curvature = (torch.tile(torch.argmax(torch.permute(batch['uv_curvature'], [0,2,3,1]), dim=-1)[:,:,:,None], [1,1,1,3]) / 6)
pred_img = torch.concatenate([torch.permute(F.interpolate(output_levels[i][0], size=(512, 256), mode='bilinear'), [0,2,3,1]) for i in range(len(output_levels))], dim=3)
uv_normals = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.4) / 2.8
pred_img = torch.concatenate([torch.permute(model.resizers[512](output_levels[i][0]), [0,2,3,1]) for i in range(len(output_levels))], dim=3)
# pred_img = (pred_img - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis]) / (pred_img.max(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis] - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis])
pred_img = (pred_img - pred_img.flatten().reshape(pred_img.shape[0], -1).min(1)[0][:,None,None,None]) / (pred_img.flatten().reshape(pred_img.shape[0], -1).max(1)[0][:,None,None,None] - pred_img.flatten().reshape(pred_img.shape[0], -1).min(1)[0][:,None,None,None])
img = torch.concat((uv_static_img, uv_body_posed_img, uv_curvature, orig_img, pred_img), axis=2)
img = torch.concat((uv_static_img, uv_body_posed_img, uv_normals, orig_img, pred_img), axis=2)
imgs = [wandb.Image(img[i].cpu().numpy()) for i in range(4)]
_F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
@@ -521,7 +548,14 @@ def plot_everything(batch: dict, type: str):
collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(gt_posed))]
figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i]) for i in range(4) }
return { **figs, **figs_coll, **{ type + '/images': imgs, } }
batch_id = batch['id']
batch_frame_nr = batch['frame_nr'].cpu().numpy()
batch_garment = batch['garment']
batch_info = [f'batch_info: id: {batch_id[i]}, frame_nr: {batch_frame_nr[i]}, garment: {batch_garment[i]}' for i in range(len(batch_id))]
batch_info = text_to_html(batch_info)
batch_info = wandb.Html(batch_info)
return { **figs, **figs_coll, **{ type + '/images': imgs, }, **{ f'{type}/batch_info': batch_info }, }
def sample(folder):
@@ -793,7 +827,7 @@ def train():
if train_loss['train/euclid_error'] < config['balancing_activation_limit']:
g_balancing_coeff += 1 / 10000
if step % 10000 == 0:
if step % 25000 == 0:
ckpt_saver(model, step, optimizer, None)
step += 1
@@ -840,12 +874,12 @@ if __name__ == '__main__':
if len(config['pyramid']['resolutions']) == 1:
for param in model.static_enc.parameters(): param.requires_grad = False
for param in model.body_enc.parameters(): param.requires_grad = False
for param in model.curv_enc.parameters(): param.requires_grad = False
for param in model.norm_enc.parameters(): param.requires_grad = False
else:
load_model(model, config['pyramid']['pretrained_pyramid'], iteration=(config['pyramid']['pretrained_pyramid_iteration'] if 'pretrained_pyramid_iteration' in config.keys() else ''))
for param in model.static_enc.parameters(): param.requires_grad = False
for param in model.body_enc.parameters(): param.requires_grad = False
for param in model.curv_enc.parameters(): param.requires_grad = False
for param in model.norm_enc.parameters(): param.requires_grad = False
# load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer'])
# model.levels[0].upscale_enhancer.trainable = False
@@ -866,6 +900,6 @@ if __name__ == '__main__':
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()