def train(modelin=args.model, modelout=args.out, device=args.device, opt=args.opt): # define model, dataloader, 3dmm eigenvectors, optimization method calib_net = PointNet(n=1) sfm_net = PointNet(n=199) if modelin != "": calib_path = os.path.join('model', 'calib_' + modelin) sfm_path = os.path.join('model', 'sfm_' + modelin) pretrained1 = torch.load(calib_path) pretrained2 = torch.load(sfm_path) calib_dict = calib_net.state_dict() sfm_dict = sfm_net.state_dict() pretrained1 = {k: v for k, v in pretrained1.items() if k in calib_dict} pretrained2 = {k: v for k, v in pretrained2.items() if k in sfm_dict} calib_dict.update(pretrained1) sfm_dict.update(pretrained2) calib_net.load_state_dict(pretrained1) sfm_net.load_state_dict(pretrained2) calib_net.to(device=device) sfm_net.to(device=device) opt1 = torch.optim.Adam(calib_net.parameters(), lr=1e-4) opt2 = torch.optim.Adam(sfm_net.parameters(), lr=1e-2) # dataloader data = dataloader.Data() loader = data.batchloader batch_size = data.batchsize # mean shape and eigenvectors for 3dmm mu_lm = torch.from_numpy(data.mu_lm).float() #.to(device=device) mu_lm[:, 2] = mu_lm[:, 2] * -1 mu_lm = torch.stack(batch_size * [mu_lm.to(device=device)]) shape = mu_lm lm_eigenvec = torch.from_numpy(data.lm_eigenvec).float().to(device=device) lm_eigenvec = torch.stack(batch_size * [lm_eigenvec]) M = data.M N = data.N # main training loop best = 10000 for epoch in itertools.count(): for j, batch in enumerate(loader): # get the input and gt values x_cam_gt = batch['x_cam_gt'].to(device=device) shape_gt = batch['x_w_gt'].to(device=device) fgt = batch['f_gt'].to(device=device) x_img = batch['x_img'].to(device=device) #beta_gt = batch['beta_gt'].to(device=device) #x_img_norm = batch['x_img_norm'] x_img_gt = batch['x_img_gt'].to(device=device).permute(0, 2, 1, 3) batch_size = fgt.shape[0] one = torch.ones(batch_size, M * N, 1).to(device=device) x_img_one = torch.cat([x_img, one], dim=2) x_cam_pt = x_cam_gt.permute(0, 1, 3, 2).reshape(batch_size, 6800, 3) x = x_img.permute(0, 2, 1) #x = x_img.permute(0,2,1).reshape(batch_size,2,M,N) ptsI = x_img_one.reshape(batch_size, M, N, 3).permute(0, 1, 3, 2)[:, :, :2, :] # if just optimizing if not opt: # calibration f = calib_net(x) + 300 K = torch.zeros((batch_size, 3, 3)).float().to(device=device) K[:, 0, 0] = f.squeeze() K[:, 1, 1] = f.squeeze() K[:, 2, 2] = 1 # sfm betas = sfm_net(x) betas = betas.unsqueeze(-1) shape = mu_lm + torch.bmm(lm_eigenvec, betas).squeeze().view( batch_size, N, 3) opt1.zero_grad() opt2.zero_grad() f_error = torch.mean(torch.abs(f - fgt)) #error2d = torch.mean(torch.abs(pred - x_img_gt)) error3d = torch.mean(torch.abs(shape - shape_gt)) error = f_error + error3d error.backward() opt1.step() opt2.step() print( f"{best:.2f} | f_error: {f_error.item():.3f} | error3d: {error3d.item():.3f} | f/fgt: {f[0].item():.1f}/{fgt[0].item():.1f} | f/fgt: {f[1].item():.1f}/{fgt[1].item():.1f} | f/fgt: {f[2].item():.1f}/{fgt[2].item():.1f} | f/fgt: {f[3].item():.1f}/{fgt[3].item():.1f} " ) continue # save model and increment weight decay torch.save(sfm_net.state_dict(), os.path.join('model', 'sfm_model.pt')) torch.save(calib_net.state_dict(), os.path.join('model', 'calib_model.pt')) ferror = test(modelin='model.pt', outfile=args.out, optimize=False) if ferror < best: best = ferror print("saving!") torch.save(sfm_net.state_dict(), os.path.join('model', 'sfm_' + modelout)) torch.save(calib_net.state_dict(), os.path.join('model', 'calib_' + modelout)) sfm_net.train() calib_net.train()
def train(modelin=args.model, modelout=args.out,device=args.device,opt=args.opt,ft=args.ft): # define model, dataloader, 3dmm eigenvectors, optimization method calib_net = PointNet(n=1,feature_transform=ft) sfm_net = PointNet(n=199,feature_transform=ft) if modelin != "": calib_path = os.path.join('model','calib_' + modelin) sfm_path = os.path.join('model','sfm_' + modelin) pretrained1 = torch.load(calib_path) pretrained2 = torch.load(sfm_path) calib_dict = calib_net.state_dict() sfm_dict = sfm_net.state_dict() pretrained1 = {k: v for k,v in pretrained1.items() if k in calib_dict} pretrained2 = {k: v for k,v in pretrained2.items() if k in sfm_dict} calib_dict.update(pretrained1) sfm_dict.update(pretrained2) calib_net.load_state_dict(pretrained1) sfm_net.load_state_dict(pretrained2) calib_net.to(device=device) sfm_net.to(device=device) opt1 = torch.optim.Adam(calib_net.parameters(),lr=1e-3) opt2 = torch.optim.Adam(sfm_net.parameters(),lr=1e-3) # dataloader loader = dataloader.SyntheticLoader() batch_size = 100 M = loader.M N = loader.N # mean shape and eigenvectors for 3dmm mu_lm = torch.from_numpy(loader.mu_lm).float()#.to(device=device) mu_lm[:,2] = mu_lm[:,2] * -1 mu_lm = torch.stack(300 * [mu_lm.to(device=device)]) shape = mu_lm lm_eigenvec = torch.from_numpy(loader.lm_eigenvec).float().to(device=device) sigma = torch.from_numpy(loader.sigma).float().detach().to(device=device) sigma = torch.diag(sigma.squeeze()) lm_eigenvec = torch.mm(lm_eigenvec, sigma) lm_eigenvec = torch.stack(300 * [lm_eigenvec]) # main training loop best = 10000 for epoch in itertools.count(): for i in range(len(loader)): if i < 3: continue v1 = loader[i] v2 = loader[i-1] v3 = loader[i-2] batch = stackVideos([v1,v2,v3],100,68,device=device) # get the input and gt values alpha_gt = batch['alpha'] x_cam_gt = batch['x_cam_gt'] shape_gt = batch['x_w_gt'] fgt = batch['f_gt'] x = batch['x_img'] M = x.shape[0] N = x.shape[-1] # calibration f = torch.squeeze(calib_net(x) + 300) K = torch.zeros((M,3,3)).float().to(device=device) K[:,0,0] = f K[:,1,1] = f K[:,2,2] = 1 # sfm alpha = sfm_net(x) alpha = alpha.unsqueeze(-1) shape = mu_lm + torch.bmm(lm_eigenvec,alpha).squeeze().view(M,N,3) shape[0:100] = shape[0:100] - shape[0:100].mean(1).unsqueeze(1) shape[100:200] = shape[100:200] - shape[100:200].mean(1).unsqueeze(1) shape[200:300] = shape[200:300] - shape[200:300].mean(1).unsqueeze(1) opt1.zero_grad() opt2.zero_grad() f1_error = torch.mean(torch.abs(f[0:100] - fgt[0])) f2_error = torch.mean(torch.abs(f[100:200] - fgt[1])) f3_error = torch.mean(torch.abs(f[200:300] - fgt[2])) #a1_error = torch.mean(torch.abs(alpha[0:100] - alpha_gt[0])) #a2_error = torch.mean(torch.abs(alpha[100:200] - alpha_gt[1])) #a3_error = torch.mean(torch.abs(alpha[200:300] - alpha_gt[2])) s1_error = torch.mean(torch.abs(shape[0:100] - shape_gt[0].unsqueeze(0))) s2_error = torch.mean(torch.abs(shape[100:200] - shape_gt[1].unsqueeze(0))) s3_error = torch.mean(torch.abs(shape[200:300] - shape_gt[2].unsqueeze(0))) ferror = f1_error + f2_error + f3_error #aerror = a1_error + a2_error + a3_error serror = s1_error + s2_error + s3_error #f_error = torch.mean(torch.abs(f - fgt)) #error3d = torch.mean(torch.norm(shape - shape_gt,dim=2)) #error = ferror + aerror error = ferror + serror error.backward() opt1.step() opt2.step() print(f"iter: {i} | best: {best:.2f} | f_error: {ferror.item():.3f} | serror: {serror.item():.3f} ") if i == 1000: break # save model and increment weight decay torch.save(sfm_net.state_dict(), os.path.join('model','sfm_model.pt')) torch.save(calib_net.state_dict(), os.path.join('model','calib_model.pt')) ferror = test(modelin='model.pt',outfile=args.out,optimize=False) if ferror < best: best = ferror print("saving!") torch.save(sfm_net.state_dict(), os.path.join('model','sfm_'+modelout)) torch.save(calib_net.state_dict(), os.path.join('model','calib_'+modelout)) sfm_net.train() calib_net.train()
def train(modelin=args.model, modelout=args.out, device=args.device, opt=args.opt, ft=args.ft): # define model, dataloader, 3dmm eigenvectors, optimization method calib_net = PointNet(n=1, feature_transform=ft) sfm_net = PointNet(n=199, feature_transform=ft) if modelin != "": calib_path = os.path.join('model', 'calib_' + modelin) sfm_path = os.path.join('model', 'sfm_' + modelin) pretrained1 = torch.load(calib_path) pretrained2 = torch.load(sfm_path) calib_dict = calib_net.state_dict() sfm_dict = sfm_net.state_dict() pretrained1 = {k: v for k, v in pretrained1.items() if k in calib_dict} pretrained2 = {k: v for k, v in pretrained2.items() if k in sfm_dict} calib_dict.update(pretrained1) sfm_dict.update(pretrained2) calib_net.load_state_dict(pretrained1) sfm_net.load_state_dict(pretrained2) calib_net.to(device=device) sfm_net.to(device=device) opt1 = torch.optim.Adam(calib_net.parameters(), lr=1e-3) opt2 = torch.optim.Adam(sfm_net.parameters(), lr=1e-3) # dataloader loader = dataloader.SyntheticLoader() batch_size = 100 M = loader.M N = loader.N # mean shape and eigenvectors for 3dmm mu_lm = torch.from_numpy(loader.mu_lm).float() #.to(device=device) mu_lm[:, 2] = mu_lm[:, 2] * -1 mu_lm = torch.stack(batch_size * [mu_lm.to(device=device)]) shape = mu_lm lm_eigenvec = torch.from_numpy( loader.lm_eigenvec).float().to(device=device) sigma = torch.from_numpy(loader.sigma).float().detach().to(device=device) sigma = torch.diag(sigma.squeeze()) lm_eigenvec = torch.mm(lm_eigenvec, sigma) lm_eigenvec = torch.stack(M * [lm_eigenvec]) # main training loop best = 10000 for epoch in itertools.count(): for j, batch in enumerate(loader): # get the input and gt values x_cam_gt = batch['x_cam_gt'].to(device=device) shape_gt = batch['x_w_gt'].to(device=device) fgt = batch['f_gt'].to(device=device) x_img = batch['x_img'].to(device=device) #beta_gt = batch['beta_gt'].to(device=device) #x_img_norm = batch['x_img_norm'] #x_img_gt = batch['x_img_gt'].to(device=device).permute(0,2,1,3) x = x_img.reshape(M, N, 2).permute(0, 2, 1) batch_size = fgt.shape[0] #x_cam_pt = x_cam_gt.permute(0,1,3,2).reshape(batch_size,6800,3) #x = x_img.permute(0,2,1).reshape(batch_size,2,M,N) #ptsI = x_img_one.reshape(batch_size,M,N,3).permute(0,1,3,2)[:,:,:2,:] # calibration f = torch.squeeze(calib_net(x) + 300) K = torch.zeros((M, 3, 3)).float().to(device=device) K[:, 0, 0] = f K[:, 1, 1] = f K[:, 2, 2] = 1 # sfm betas = sfm_net(x) betas = betas.unsqueeze(-1) shape = mu_lm + torch.bmm(lm_eigenvec, betas).squeeze().view( M, N, 3) shape = shape - shape.mean(1).unsqueeze(1) opt1.zero_grad() opt2.zero_grad() f_error = torch.mean(torch.abs(f - fgt)) #error2d = torch.mean(torch.abs(pred - x_img_gt)) error3d = torch.mean(torch.norm(shape - shape_gt, dim=2)) error = f_error + error3d error.backward() opt1.step() opt2.step() print( f"iter: {j} | best: {best:.2f} | f_error: {f_error.item():.3f} | error3d: {error3d.item():.3f} " ) if j == 1000: break # save model and increment weight decay torch.save(sfm_net.state_dict(), os.path.join('model', 'sfm_model.pt')) torch.save(calib_net.state_dict(), os.path.join('model', 'calib_model.pt')) ferror = test(modelin='model.pt', outfile=args.out, optimize=False) if ferror < best: best = ferror print("saving!") torch.save(sfm_net.state_dict(), os.path.join('model', 'sfm_' + modelout)) torch.save(calib_net.state_dict(), os.path.join('model', 'calib_' + modelout)) sfm_net.train() calib_net.train()