def test(): end = time.time() batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() losses1 = AverageMeter() iou_all = [] for class_id, class_name in dataset_val.class_ids_pair: directory_mesh_cls = os.path.join(directory_mesh, class_id) os.makedirs(directory_mesh_cls, exist_ok=True) iou = 0 for i, (im, vx) in enumerate( dataset_val.get_all_batches_for_evaluation( args.batch_size, class_id)): images = torch.autograd.Variable(im).cuda() voxels = vx.numpy() batch_iou, vertices, faces = model(images, voxels=voxels, task='test') iou += batch_iou.sum() batch_time.update(time.time() - end) end = time.time() # save demo images for k in range(vertices.size(0)): obj_id = (i * args.batch_size + k) if obj_id % args.save_freq == 0: mesh_path = os.path.join(directory_mesh_cls, '%06d.obj' % obj_id) input_path = os.path.join(directory_mesh_cls, '%06d.png' % obj_id) srf.save_obj(mesh_path, vertices[k], faces[k]) imageio.imsave(input_path, img_cvt(images[k])) # print loss if i % args.print_freq == 0: print('Iter: [{0}/{1}]\t' 'Time {batch_time.val:.3f}\t' 'IoU {2:.3f}\t'.format( i, ((dataset_val.num_data[class_id] * 24) // args.batch_size), batch_iou.mean(), batch_time=batch_time)) iou_cls = iou / 24. / dataset_val.num_data[class_id] * 100 iou_all.append(iou_cls) print('=================================') print('Mean IoU: %.3f for class %s' % (iou_cls, class_name)) print('\n') print('=================================') print('Mean IoU: %.3f for all classes' % (sum(iou_all) / len(iou_all)))
def train(): writer = SummaryWriter() print("Made summary writer") end = time.time() batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() for i in range(start_iter, args.num_iterations + 1): # adjust learning rate and sigma_val (decay after 150k iter) lr = adjust_learning_rate([optimizer], args.learning_rate, i, method=args.lr_type) model.set_sigma(adjust_sigma(args.sigma_val, i)) # TRAIN # load images from multi-view images_a, images_b, viewpoints_a, viewpoints_b = dataset_train.get_random_batch( args.batch_size) images_a = images_a.cuda() images_b = images_b.cuda() viewpoints_a = viewpoints_a.cuda() viewpoints_b = viewpoints_b.cuda() # soft render images render_images, laplacian_loss, flatten_loss = model( [images_a, images_b], [viewpoints_a, viewpoints_b], task='train') laplacian_loss = laplacian_loss.mean() flatten_loss = flatten_loss.mean() # compute loss loss = multiview_iou_loss(render_images, images_a, images_b) + \ args.lambda_laplacian * laplacian_loss + \ args.lambda_flatten * flatten_loss # VAL # load images from multi-view val_images_a, val_images_b, val_viewpoints_a, val_viewpoints_b = dataset_val.get_random_batch( args.batch_size) val_images_a = val_images_a.cuda() val_images_b = val_images_b.cuda() val_viewpoints_a = val_viewpoints_a.cuda() val_viewpoints_b = val_viewpoints_b.cuda() # soft render images val_render_images, val_laplacian_loss, val_flatten_loss = model( [val_images_a, val_images_b], [val_viewpoints_a, val_viewpoints_b], task='train') val_laplacian_loss = val_laplacian_loss.mean() val_flatten_loss = val_flatten_loss.mean() # compute loss val_loss = multiview_iou_loss(val_render_images, val_images_a, val_images_b) + \ args.lambda_laplacian * val_laplacian_loss + \ args.lambda_flatten * val_flatten_loss # RECORD LOSSES writer.add_scalar('Loss/train', loss.data.item(), i) writer.add_scalar('Loss/val', val_loss.data.item(), i) losses.update(loss.data.item(), images_a.size(0)) # compute gradient and optimize optimizer.zero_grad() loss.backward() optimizer.step() batch_time.update(time.time() - end) end = time.time() # save checkpoint if i % args.save_freq == 0: model_path = os.path.join(directory_output, 'checkpoint_%07d.pth.tar' % i) torch.save( { 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), }, model_path) # save demo images if i % args.demo_freq == 0: demo_image = images_a[0:1] demo_path = os.path.join(directory_output, 'demo_%07d.obj' % i) demo_v, demo_f = model.reconstruct(demo_image) srf.save_obj(demo_path, demo_v[0], demo_f[0]) fake_img = img_cvt(render_images[0][0]) real_img = img_cvt(images_a[0]) imageio.imsave(os.path.join(image_output, '%07d_fake.png' % i), fake_img) imageio.imsave(os.path.join(image_output, '%07d_input.png' % i), real_img) # writer.add_image('Fake image', fake_img, i) # writer.add_image('Real image', real_img, i) # print if i % args.print_freq == 0: print('Iter: [{0}/{1}]\t' 'Time {batch_time.val:.3f}\t' 'Loss {loss.val:.3f}\t' 'lr {lr:.6f}\t' 'sv {sv:.6f}\t'.format( i, args.num_iterations, batch_time=batch_time, loss=losses, lr=lr, sv=model.renderer.rasterizer.sigma_val)) writer.close()
def train(): end = time.time() batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() iou_losses = [] for i in range(start_iter, args.num_iterations + 1): # adjust learning rate and sigma_val (decay after 150k iter) lr = adjust_learning_rate([optimizer], args.learning_rate, i, method=args.lr_type) model.set_sigma(adjust_sigma(args.sigma_val, i)) # load images from multi-view images_a, images_b, viewpoints_a, viewpoints_b = dataset_train.get_random_batch( args.batch_size) images_a = images_a.cuda() images_b = images_b.cuda() viewpoints_a = viewpoints_a.cuda() viewpoints_b = viewpoints_b.cuda() # soft render images render_images, laplacian_loss, flatten_loss = model( [images_a, images_b], [viewpoints_a, viewpoints_b], task='train') laplacian_loss = laplacian_loss.mean() flatten_loss = flatten_loss.mean() iou_loss = multiview_iou_loss(render_images, images_a, images_b) iou_losses.append(iou_loss.data.item()) # compute loss loss = iou_loss + \ args.lambda_laplacian * laplacian_loss + \ args.lambda_flatten * flatten_loss losses.update(loss.data.item(), images_a.size(0)) # compute gradient and optimize optimizer.zero_grad() loss.backward() optimizer.step() batch_time.update(time.time() - end) end = time.time() # save checkpoint if i % args.save_freq == 0: model_path = os.path.join(directory_output, 'checkpoint_%07d.pth.tar' % i) torch.save( { 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), }, model_path) # save demo images if i % args.demo_freq == 0: demo_image = images_a[0:1] demo_path = os.path.join(directory_output, 'demo_%07d.obj' % i) demo_v, demo_f = model.reconstruct(demo_image) srf.save_obj(demo_path, demo_v[0], demo_f[0]) imageio.imsave(os.path.join(image_output, '%07d_fake.png' % i), img_cvt(render_images[0][0])) imageio.imsave(os.path.join(image_output, '%07d_input.png' % i), img_cvt(images_a[0])) # print if i % args.print_freq == 0: print('Iter: [{0}/{1}]\t' 'Time {batch_time.val:.3f}\t' 'Loss {loss.val:.3f}\t' 'lr {lr:.6f}\t' 'sv {sv:.6f}\t'.format( i, args.num_iterations, batch_time=batch_time, loss=losses, lr=lr, sv=model.renderer.rasterizer.sigma_val)) print(iou_losses) with open('loss.txt', 'w') as f: f.write(str(iou_losses))
def train(): end = time.time() batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() env_name = args.experiment_id img_a_logger = VisdomLogger('images', env=env_name, port=8097, opts=dict(title='img_a')) img_b_logger = VisdomLogger('images', env=env_name, port=8097, opts=dict(title='img_b')) rnd_logger = [VisdomLogger('images', env=env_name, port=8097, opts=dict(title='rnd_rgb_{}'.format(_i))) for _i in range(4)] for i in range(start_iter, args.num_iterations + 1): # adjust learning rate and sigma_val (decay after 150k iter) lr = adjust_learning_rate([optimizer], args.learning_rate, i, method=args.lr_type) model.set_sigma(adjust_sigma(args.sigma_val, i)) # load images from multi-view images_a, images_b, viewpoints_a, viewpoints_b = dataset_train.get_random_batch(args.batch_size) images_a = images_a.cuda() images_b = images_b.cuda() viewpoints_a = viewpoints_a.cuda() viewpoints_b = viewpoints_b.cuda() # soft render images render_images, laplacian_loss, flatten_loss = model([images_a, images_b], [viewpoints_a, viewpoints_b], task='train') img_a_logger.log(F.interpolate(images_a.detach().clone(), scale_factor=4).squeeze()[:,:3]) img_b_logger.log(F.interpolate(images_b.detach().clone(), scale_factor=4).squeeze()[:,:3]) for _i in range(4): rnd_logger[_i].log(F.interpolate(render_images[_i].detach().clone(), scale_factor=4).squeeze()[:,:3]) import ipdb; ipdb.set_trace() laplacian_loss = laplacian_loss.mean() flatten_loss = flatten_loss.mean() # compute loss loss = multiview_iou_loss(render_images, images_a, images_b) + \ args.lambda_laplacian * laplacian_loss + \ args.lambda_flatten * flatten_loss losses.update(loss.data.item(), images_a.size(0)) # compute gradient and optimize optimizer.zero_grad() loss.backward() optimizer.step() batch_time.update(time.time() - end) end = time.time() # save checkpoint if i % args.save_freq == 0: model_path = os.path.join(directory_output, 'checkpoint_%07d.pth.tar'%i) torch.save({ 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), }, model_path) # save demo images if i % args.demo_freq == 0: demo_image = images_a[0:1] demo_path = os.path.join(directory_output, 'demo_%07d.obj'%i) demo_v, demo_f = model.reconstruct(demo_image) srf.save_obj(demo_path, demo_v[0], demo_f[0]) imageio.imsave(os.path.join(image_output, '%07d_fake.png' % i), img_cvt(render_images[0][0])) imageio.imsave(os.path.join(image_output, '%07d_input.png' % i), img_cvt(images_a[0])) # print if i % args.print_freq == 0: print('Iter: [{0}/{1}]\t' 'Time {batch_time.val:.3f}\t' 'Loss {loss.val:.3f}\t' 'lr {lr:.6f}\t' 'sv {sv:.6f}\t'.format(i, args.num_iterations, batch_time=batch_time, loss=losses, lr=lr, sv=model.renderer.rasterizer.sigma_val))
def train(): end = time.time() batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() run_times = -1 cnt = 0 jt.sync_all(True) sta = time.time() for i in range(start_iter, args.num_iterations + 1): # adjust learning rate and sigma_val (decay after 150k iter) lr = adjust_learning_rate([optimizer], args.learning_rate, i, method=args.lr_type) model.set_sigma(adjust_sigma(args.sigma_val, i)) # load images from multi-view images_a, images_b, viewpoints_a, viewpoints_b = dataset_train.get_random_batch( args.batch_size) # soft render images render_images, laplacian_loss, flatten_loss = model( [images_a, images_b], [viewpoints_a, viewpoints_b], task='train') laplacian_loss = laplacian_loss.mean() flatten_loss = flatten_loss.mean() # compute loss loss = multiview_iou_loss(render_images, images_a, images_b) + \ args.lambda_laplacian * laplacian_loss + \ args.lambda_flatten * flatten_loss losses.update(loss.data, images_a.size(0)) # compute gradient and optimize optimizer.step(loss) batch_time.update(time.time() - end) end = time.time() if run_times != -1: if cnt == run_times: jt.sync_all(True) print( f"Costs {time.time() - sta} secs for running {run_times} times." ) exit(0) else: cnt += 1 continue # save checkpoint if i % args.save_freq == 0: model_path = os.path.join(directory_output, 'checkpoint_%07d.pkl' % i) model.save(model_path) # save demo images if i % args.demo_freq == 0: demo_image = images_a[0:1] demo_path = os.path.join(directory_output, 'demo_%07d.obj' % i) demo_v, demo_f = model.reconstruct(demo_image) jr.save_obj(demo_path, demo_v[0], demo_f[0]) imageio.imsave(os.path.join(image_output, '%07d_fake.png' % i), img_cvt(render_images[0][0][None, :, :])) imageio.imsave(os.path.join(image_output, '%07d_input.png' % i), img_cvt(images_a[0])) # print if i % args.print_freq == 0: print('Iter: [{0}/{1}]\t' 'Time {batch_time:.3f}\t' 'Loss {loss:.3f}\t' 'lr {lr:.6f}\t' 'sv 1\t'.format(i, args.num_iterations, batch_time=batch_time.val, loss=losses.val.tolist()[0], lr=lr, sv=model.renderer.rasterizer.sigma_val))