def validate(epoch, val_loader, model, crit_cls, crit_reg, opt, ctx, gen_shape=False): """ One validation """ generated_shapes = [] original_shapes = [] sample_prob = opt.inner_sample_prob loss_cls_sum, loss_reg_sum, n = 0.0, 0.0, 0 for idx, data in enumerate(val_loader): start = time.time() shapes, labels, masks, params, param_masks = data[0], data[1], data[ 2], data[3], data[4] gt = shapes shapes = nd.expand_dims(shapes, axis=1) shapes = shapes.as_in_context(ctx) labels = labels.as_in_context(ctx) masks = masks.as_in_context(ctx) params = params.as_in_context(ctx) param_masks = param_masks.as_in_context(ctx) with autograd.train_mode(): out = model.decode(shapes) #out = model(shapes, labels, sample_prob) bsz, n_block, n_step = labels.shape labels = labels.reshape(bsz, n_block * n_step) masks = masks.reshape(bsz, n_block * n_step) out_pgm = out[0].reshape(bsz, n_block * n_step, opt.program_size + 1) bsz, n_block, n_step, n_param = params.shape params = params.reshape(bsz, n_block * n_step, n_param) param_masks = param_masks.reshape(bsz, n_block * n_step, n_param) out_param = out[1].reshape(bsz, n_block * n_step, n_param) loss_cls, acc = crit_cls(out_pgm, labels, masks) loss_reg = crit_reg(out_param, params, param_masks) end = time.time() loss_cls = loss_cls.mean().asscalar() loss_reg = loss_reg.mean().asscalar() if idx % opt.info_interval == 0: out_1 = nd.round(out[0]).astype('int64') out_2 = nd.round(out[1]).astype('int64') pred = nd.from_numpy(decode_multiple_block( out_1, out_2)).astype("float32").as_in_context(mx.cpu()) IoU = BatchIoU(pred, gt) print( "Test: epoch {} batch {}/{}, loss_cls = {:.3f}, loss_reg = {:.3f}, acc = {:.3f}, IoU = {:.3f} time = {:.3f}" .format(epoch, idx, len(val_loader), loss_cls, loss_reg, acc[0].asscalar(), IoU.mean(), end - start)) sys.stdout.flush()
def run(): opt = parse_argument() if not os.path.isdir(opt.prog_save_path): os.makedirs(opt.prog_save_path) if not os.path.isdir(opt.imgs_save_path): os.makedirs(opt.imgs_save_path) print('========= arguments =========') for key, val in vars(opt).items(): print("{:20} {}".format(key, val)) print('========= arguments =========') # data loader test_set = ShapeNet3D(opt.data) test_loader = DataLoader( dataset=test_set, batch_size=opt.batch_size, shuffle=False, num_workers=opt.num_workers, ) # model ckpt = torch.load(opt.model) model = BlockOuterNet(ckpt['opt']) model.load_state_dict(ckpt['model']) if opt.is_cuda: model = model.cuda() cudnn.benchmark = True # test the model and evaluate the IoU ori_shapes, gen_shapes, pgms, params = test_on_shapenet_data( epoch=0, test_loader=test_loader, model=model, opt=opt, gen_shape=True) IoU = BatchIoU(ori_shapes, gen_shapes) print("Mean IoU: {:.3f}".format(IoU.mean())) # execute the generated program to generate the reconstructed shapes # for double-check purpose, can be disabled num_shapes = gen_shapes.shape[0] res = [] for i in range(num_shapes): data = execute_shape_program(pgms[i], params[i]) res.append(data.reshape((1, 32, 32, 32))) res = np.concatenate(res, axis=0) IoU_2 = BatchIoU(ori_shapes, res) assert abs(IoU.mean() - IoU_2.mean()) < 0.1, 'IoUs are not matched' # save results save_file = os.path.join(opt.save_path, 'shapes.h5') f = h5py.File(save_file, 'w') f['data'] = gen_shapes f['pgms'] = pgms f['params'] = params f.close() # Interpreting programs to understandable program strings if opt.save_prog: interpreter = Interpreter(translate, rotate, end) num_programs = gen_shapes.shape[0] for i in range(min(num_programs, opt.num_render)): program = interpreter.interpret(pgms[i], params[i]) save_file = os.path.join(opt.prog_save_path, '{}.txt'.format(i)) with open(save_file, 'w') as out: out.write(program) # Visualization if opt.save_img: data = gen_shapes.transpose((0, 3, 2, 1)) data = np.flip(data, axis=2) num_shapes = data.shape[0] for i in range(min(num_shapes, opt.num_render)): voxels = data[i] save_name = os.path.join(opt.imgs_save_path, '{}.png'.format(i)) visualization(voxels, threshold=0.1, save_name=save_name, uniform_size=0.9)
def train(epoch, train_loader, model, logsoft, soft, criterion, optimizer, opt): """ one epoch training for program executor """ model.train() criterion.train() for idx, data in enumerate(train_loader): start_t = time.time() optimizer.zero_grad() shape, label, param = data[0], data[1], data[2] bsz = shape.size(0) n_step = label.size(1) index = np.array(list(map(lambda x: n_step, label))) index = index - 1 # add noise during training, making the executor accept # continuous output from program generator label = label.view(-1, 1) pgm_vector = 0.1 * torch.rand(bsz * n_step, stop_id) pgm_noise = 0.1 * torch.rand(bsz * n_step, 1) pgm_value = torch.ones(bsz * n_step, 1) - pgm_noise pgm_vector.scatter_(1, label, pgm_value) pgm_vector = pgm_vector.view(bsz, n_step, stop_id) param_noise = torch.rand(param.size()) param_vector = param + 0.6 * (param_noise - 0.5) gt = shape index = torch.from_numpy(index).long() pgm_vector = pgm_vector.float() param_vector = param_vector.float() if opt.is_cuda: gt = gt.cuda() index = index.cuda() pgm_vector = pgm_vector.cuda() param_vector = param_vector.cuda() pred = model(pgm_vector, param_vector, index) scores = logsoft(pred) loss = criterion(scores, gt) loss.backward() clip_gradient(optimizer, opt.grad_clip) optimizer.step() loss = loss.data[0] pred = soft(pred) pred = pred[:, 1, :, :, :] s1 = gt.view(-1, 32, 32, 32).data.cpu().numpy() s2 = pred.squeeze().data.cpu().numpy() s2 = (s2 > 0.5) batch_iou = BatchIoU(s1, s2) iou = batch_iou.sum() / s1.shape[0] end_t = time.time() if idx % (opt.info_interval * 10) == 0: print( "Train: epoch {} batch {}/{}, loss13 = {:.3f}, iou = {:.3f}, time = {:.3f}" .format(epoch, idx, len(train_loader), loss, iou, end_t - start_t)) sys.stdout.flush()
def run(): opt = options_train_executor.parse() print('===== arguments: program executor =====') for key, val in vars(opt).items(): print("{:20} {}".format(key, val)) print('===== arguments: program executor =====') if not os.path.isdir(opt.save_folder): os.makedirs(opt.save_folder) # build dataloader train_set = PartPrimitive(opt.train_file) train_loader = DataLoader( dataset=train_set, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers, ) val_set = PartPrimitive(opt.val_file) val_loader = DataLoader( dataset=val_set, batch_size=opt.batch_size, shuffle=False, num_workers=opt.num_workers, ) # build the model model = RenderNet(opt) logsoft = nn.LogSoftmax(dim=1) soft = nn.Softmax(dim=1) criterion = nn.NLLLoss(weight=torch.Tensor([opt.n_weight, opt.p_weight])) if opt.is_cuda: if opt.num_gpu > 1: gpu_ids = [i for i in range(opt.num_gpu)] model = torch.nn.DataParallel(model, device_ids=gpu_ids) model = model.cuda() logsoft = logsoft.cuda() soft = soft.cuda() criterion = criterion.cuda() cudnn.benchmark = True optimizer = optim.Adam(model.parameters(), lr=opt.learning_rate, betas=(opt.beta1, opt.beta2), weight_decay=opt.weight_decay) for epoch in range(1, opt.epochs + 1): adjust_learning_rate(epoch, opt, optimizer) print("###################") print("training") train(epoch, train_loader, model, logsoft, soft, criterion, optimizer, opt) print("###################") print("testing") gen_shapes, ori_shapes = validate(epoch, val_loader, model, logsoft, soft, criterion, opt, gen_shape=True) gen_shapes = (gen_shapes > 0.5) gen_shapes = gen_shapes.astype(np.float32) iou = BatchIoU(ori_shapes, gen_shapes) print("Mean IoU: {:.3f}".format(iou.mean())) if epoch % opt.save_interval == 0: print('Saving...') state = { 'opt': opt, 'model': model.module.state_dict() if opt.num_gpu > 1 else model.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch, } save_file = os.path.join( opt.save_folder, 'ckpt_epoch_{epoch}.t7'.format(epoch=epoch)) torch.save(state, save_file) state = { 'opt': opt, 'model': model.module.state_dict() if opt.num_gpu > 1 else model.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': opt.epochs, } save_file = os.path.join(opt.save_folder, 'program_executor.t7') torch.save(state, save_file)
def validate(epoch, val_loader, model, logsoft, soft, criterion, opt, gen_shape=False): # load pre-fixed randomization try: rand1 = np.load(opt.rand1) rand2 = np.load(opt.rand2) rand3 = np.load(opt.rand3) except: rand1 = np.random.rand(opt.batch_size * opt.seq_length, stop_id).astype(np.float32) rand2 = np.random.rand(opt.batch_size * opt.seq_length, 1).astype(np.float32) rand3 = np.random.rand(opt.batch_size, opt.seq_length, max_param - 1).astype(np.float32) np.save(opt.rand1, rand1) np.save(opt.rand2, rand2) np.save(opt.rand3, rand3) model.eval() criterion.eval() generated_shapes = [] original_shapes = [] for idx, data in enumerate(val_loader): start_t = time.time() shape, label, param = data[0], data[1], data[2] bsz = shape.size(0) n_step = label.size(1) index = np.array(list(map(lambda x: n_step, label))) index = index - 1 label = label.view(-1, 1) pgm_vector = 0.1 * torch.from_numpy(rand1) pgm_noise = 0.1 * torch.from_numpy(rand2) pgm_value = torch.ones(bsz * n_step, 1) - pgm_noise pgm_vector.scatter_(1, label, pgm_value) pgm_vector = pgm_vector.view(bsz, n_step, stop_id) param_noise = torch.from_numpy(rand3) param_vector = param + 0.6 * (param_noise - 0.5) gt = shape index = torch.from_numpy(index).long() pgm_vector = pgm_vector.float() param_vector = param_vector.float() if opt.is_cuda: gt = gt.cuda() index = index.cuda() pgm_vector = pgm_vector.cuda() param_vector = param_vector.cuda() pred = model(pgm_vector, param_vector, index) scores = logsoft(pred) loss = criterion(scores, gt) loss = loss.data[0] pred = soft(pred) pred = pred[:, 1, :, :, :] s1 = gt.view(-1, 32, 32, 32).data.cpu().numpy() s2 = pred.squeeze().data.cpu().numpy() s2 = (s2 > 0.5) batch_iou = BatchIoU(s1, s2) iou = batch_iou.sum() / s1.shape[0] original_shapes.append(s1) generated_shapes.append(s2) end_t = time.time() if (idx + 1) % opt.info_interval == 0: print( "Test: epoch {} batch {}/{}, loss13 = {:.3f}, iou = {:.3f}, time = {:.3f}" .format(epoch, idx + 1, len(val_loader), loss, iou, end_t - start_t)) sys.stdout.flush() if gen_shape: generated_shapes = np.asarray(generated_shapes) original_shapes = np.asarray(original_shapes) return generated_shapes, original_shapes
def train(epoch, train_loader, model,loss, optimizer, opt,ctx,train_loss,train_iou): """ one epoch training for program executor """ loss_sum,iou_sum,n = 0.0,0.0,0 for idx, data in enumerate(train_loader): start_t = time.time() shape, label, param = data bsz = shape.shape[0] n_step = label.shape[1] #print("label.shape:",label) #print("n_step:",n_step,"bsz:",bsz,"stop_id:",stop_id) index = np.array(list(map(lambda x: n_step, label)))-1 #index = label # add noise during training, making the executor accept # continuous output from program generator label = label.reshape(-1,1).asnumpy() pgm_vector = 0.2 * np.random.uniform(0,1,(bsz * n_step, stop_id)) pgm_noise = 0.2 *np.random.uniform(0,1,label.shape) pgm_value = 1 - pgm_noise #print('pgm_val.shape:',pgm_value.shape,'label.shape:',label.shape,'label.shape:',label.shape) pgm_vector = scatter_numpy(pgm_vector,1,label,pgm_value).reshape(bsz,n_step,stop_id) param_noise = nd.random_uniform(0,1,shape=param.shape) param_vector = param + 0.6 * (param_noise - 0.5) #print("param_vector.shape:",param_vector.shape) gt = shape.as_in_context(ctx) #print(pgm_vector.dtype) index = nd.from_numpy(index).astype('int64').as_in_context(ctx) pgm_vector = nd.from_numpy(pgm_vector).astype('float32').as_in_context(ctx) param_vector = param_vector.as_in_context(ctx) with autograd.record(): pred = model(pgm_vector, param_vector, index) scores = nd.log_softmax(pred,axis=1) pred0 = scores[:,0].squeeze()*opt.n_weight pred1 = scores[:,1].squeeze()*opt.p_weight l = -nd.where(gt, pred1, pred0).mean((1,2,3)) #l = -(nd.pick(scores1, gt, axis=1, keepdims=True)*opt.n_weight # +nd.pick(scores2,(1-gt), axis=1, keepdims=True)*opt.p_weight).mean((1,2,3,4)) l.backward() #clip_gradient(optimizer, opt.grad_clip) #optimizer._allreduce_grads(); optimizer.step(l.shape[0],ignore_stale_grad=True) l = l.mean().asscalar() pred = nd.softmax(pred,axis = 1) pred = pred[:, 1, :, :, :] s1 = gt.reshape(-1, 32, 32, 32).astype('float32').as_in_context(mx.cpu()) s2 = pred.squeeze().as_in_context(mx.cpu()) #print(s2.shape) s2 = (s2 > 0.5) batch_iou = BatchIoU(s1, s2) iou = batch_iou.mean() end_t = time.time() loss_sum+=l n+=1 iou_sum+=iou if idx % (opt.info_interval * 10) == 0: print("Train: epoch {} batch {}/{}, loss13 = {:.3f}, iou = {:.3f}, time = {:.3f}" .format(epoch, idx, len(train_loader), l, iou, end_t - start_t)) sys.stdout.flush() train_loss.append(loss_sum/n) train_iou.append(iou_sum/n)
def validate(epoch, val_loader, model, loss, opt, ctx,val_loss,val_iou, gen_shape=False): # load pre-fixed randomization try: rand1 = np.load(opt.rand1) rand2 = np.load(opt.rand2) rand3 = np.load(opt.rand3) except: rand1 = np.random.rand(opt.batch_size * opt.seq_length, stop_id).astype(np.float32) rand2 = np.random.rand(opt.batch_size * opt.seq_length, 1).astype(np.float32) rand3 = np.random.rand(opt.batch_size, opt.seq_length, max_param - 1).astype(np.float32) np.save(opt.rand1, rand1) np.save(opt.rand2, rand2) np.save(opt.rand3, rand3) generated_shapes = None original_shapes = None loss_sum,iou_sum,n = 0.0,0.0,0 for idx, data in enumerate(val_loader): start_t = time.time() shape, label, param = data bsz = shape.shape[0] n_step = label.shape[1] index = np.array(list(map(lambda x: n_step, label))) index = index - 1 # add noise during training, making the executor accept # continuous output from program generator label = label.reshape(-1,1).asnumpy() pgm_vector = 0.1*rand1 pgm_noise = 0.1*rand2 pgm_value = np.ones(label.shape) - pgm_noise #print('pgm_val.shape:',pgm_value.shape,'label.shape:',label.shape,'label.shape:',label.shape) pgm_vector = scatter_numpy(pgm_vector,1,label,pgm_value).reshape(bsz,n_step,stop_id) param_noise = nd.from_numpy(rand3) #print(param.shape,param_noise.shape) param_vector = param + 0.6 * (param_noise - 0.5) gt = shape.astype('float32').as_in_context(ctx) index = nd.from_numpy(index).astype('int64').as_in_context(ctx) pgm_vector = nd.from_numpy(pgm_vector).as_in_context(ctx) param_vector = param_vector.as_in_context(ctx) #prediction pred = model(pgm_vector, param_vector, index) scores = nd.log_softmax(pred,axis=1) pred0 = scores[:,0].squeeze()*opt.p_weight pred1 = scores[:,1].squeeze()*opt.n_weight l = -nd.where(gt, pred1, pred0).mean((1,2,3)) #print(pred2.dtype,gt.dtype) #l = loss(pred,gt,sample_weight = nd.array([opt.n_weight,opt.p_weight])) l = l.mean().asscalar() pred = nd.softmax(pred,axis=1) pred = pred[:, 1, :, :, :] s1 = gt.reshape(-1, 32, 32, 32).as_in_context(mx.cpu()) s2 = pred.squeeze().as_in_context(mx.cpu()) s2 = (s2 > 0.5) batch_iou = BatchIoU(s1, s2) iou = batch_iou.mean() loss_sum+=l n+=1 iou_sum+=iou if(idx+1)%5==0 and gen_shape: if original_shapes is None: original_shapes = s1.expand_dims(axis=0) generated_shapes = s2.expand_dims(axis=0) else: original_shapes = nd.concat(original_shapes,s1.expand_dims(axis=0),dim=0) generated_shapes = nd.concat(generated_shapes,s2.expand_dims(axis=0),dim=0) end_t = time.time() if (idx + 1) % opt.info_interval == 0: print("Test: epoch {} batch {}/{}, loss13 = {:.3f}, iou = {:.3f}, time = {:.3f}" .format(epoch, idx + 1, len(val_loader), l, iou, end_t - start_t)) sys.stdout.flush() if(idx+1>len(val_loader)/10): break; val_loss.append(loss_sum/n) val_iou.append(iou_sum/n) return generated_shapes, original_shapes
def train(epoch, train_loader, generator, executor, soft, criterion, optimizer, opt): """ one epoch guided adaptation """ generator.train() # set executor as train, but actually does not update parameters # otherwise cannot bp through LSTM executor.train() def set_bn_eval(m): classname = m.__class__.__name__ if classname.find('BatchNorm') != -1: m.eval() executor.apply(set_bn_eval) for idx, data in enumerate(train_loader): start = time.time() optimizer.zero_grad() generator.zero_grad() executor.zero_grad() shapes = data raw_shapes = data shapes = torch.unsqueeze(shapes, 1) if opt.is_cuda: shapes = shapes.cuda() pgms, params = generator.decode(shapes) # truly rendered shapes rendered_shapes = decode_multiple_block(pgms, params) IoU2 = BatchIoU(rendered_shapes, raw_shapes.clone().numpy()) # neurally rendered shapes pgms = torch.exp(pgms) bsz, n_block, n_step, n_vocab = pgms.shape pgm_vector = pgms.view(bsz * n_block, n_step, n_vocab) bsz, n_block, n_step, n_param = params.shape param_vector = params.view(bsz * n_block, n_step, n_param) index = (n_step - 1) * torch.ones(bsz * n_block).long() if opt.is_cuda: index = index.cuda() pred = executor(pgm_vector, param_vector, index) pred = soft(pred) pred = pred[:, 1, :, :, :] pred = pred.contiguous().view(bsz, n_block, 32, 32, 32) rec, _ = torch.max(pred[:, :, :, :, :], dim=1) rec1 = rec rec1.unsqueeze_(1) rec0 = 1 - rec1 rec_all = torch.cat((rec0, rec1), dim=1) rec_all = torch.log(rec_all + 1e-10) loss = criterion(rec_all, shapes.detach().squeeze_(1).long()) loss.backward() clip_gradient(optimizer, opt.grad_clip) optimizer.step() reconstruction = rec.data.cpu().numpy() reconstruction = np.squeeze(reconstruction, 1) reconstruction = reconstruction > 0.5 reconstruction = reconstruction.astype(np.uint8) raw_shapes = raw_shapes.clone().numpy() IoU1 = BatchIoU(reconstruction, raw_shapes) if opt.is_cuda: torch.cuda.synchronize() end = time.time() if idx % opt.info_interval == 0: print( "Train: epoch {} batch {}/{}, loss = {:.3f}, IoU1 = {:.3f}, IoU2 = {:.3f}, time = {:.3f}" .format(epoch, idx, len(train_loader), loss.data[0], IoU1.mean(), IoU2.mean(), end - start)) sys.stdout.flush()
def run(): # get options opt = options_guided_adaptation.parse() print('===== arguments: guided adaptation =====') for key, val in vars(opt).items(): print("{:20} {}".format(key, val)) print('===== arguments: guided adaptation =====') if not os.path.isdir(opt.save_folder): os.makedirs(opt.save_folder) # build loaders train_set = ShapeNet3D(opt.train_file) train_loader = DataLoader( dataset=train_set, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers, ) val_set = ShapeNet3D(opt.val_file) val_loader = DataLoader( dataset=val_set, batch_size=opt.batch_size, shuffle=False, num_workers=opt.num_workers, ) # load program generator ckpt_p_gen = torch.load(opt.p_gen_path) generator = BlockOuterNet(ckpt_p_gen['opt']) generator.load_state_dict(ckpt_p_gen['model']) # load program executor ckpt_p_exe = torch.load(opt.p_exe_path) executor = RenderNet(ckpt_p_exe['opt']) executor.load_state_dict(ckpt_p_exe['model']) # build loss functions soft = nn.Softmax(dim=1) criterion = nn.NLLLoss(weight=torch.Tensor([1, 1])) if opt.is_cuda: generator = generator.cuda() executor = executor.cuda() soft = soft.cuda() criterion = criterion.cuda() cudnn.benchmark = True optimizer = optim.Adam(generator.parameters(), lr=opt.learning_rate, betas=(opt.beta1, opt.beta2), weight_decay=opt.weight_decay) print("###################") print("testing") gen_shapes, ori_shapes = validate(0, val_loader, generator, opt, gen_shape=True) IoU = BatchIoU(ori_shapes, gen_shapes) print("iou: ", IoU.mean()) best_iou = 0 for epoch in range(1, opt.epochs + 1): print("###################") print("adaptation") train(epoch, train_loader, generator, executor, soft, criterion, optimizer, opt) print("###################") print("testing") gen_shapes, ori_shapes = validate(epoch, val_loader, generator, opt, gen_shape=True) IoU = BatchIoU(ori_shapes, gen_shapes) print("iou: ", IoU.mean()) if epoch % opt.save_interval == 0: print('Saving...') state = { 'opt': ckpt_p_gen['opt'], 'model': generator.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch, } save_file = os.path.join( opt.save_folder, 'ckpt_epoch_{epoch}.t7'.format(epoch=epoch)) torch.save(state, save_file) if IoU.mean() >= best_iou: print('Saving best model') state = { 'opt': ckpt_p_gen['opt'], 'model': generator.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch, } save_file = os.path.join( opt.save_folder, 'program_generator_GA_{}.t7'.format(opt.cls)) torch.save(state, save_file) best_iou = IoU.mean()
def train(epoch, train_loader, generator, executor, criterion, optimizer, opt, ctx): """ one epoch guided adaptation """ def set_bn_eval(m): if m.prefix[:9] == 'batchnorm': m._kwargs['use_global_stats'] = True m.grad_req = 'null' executor.apply(set_bn_eval) for idx, data in enumerate(train_loader): start = time.time() shapes = data.as_in_context(ctx) raw_shapes = data shapes = shapes.expand_dims(axis=1) with autograd.record(): pgms, params = generator.decode(shapes) # truly rendered shapes pgms_int = nd.round(pgms).astype('int64') params_int = nd.round(params).astype('int64') # neurally rendered shapes pgms = nd.exp(pgms) bsz, n_block, n_step, n_vocab = pgms.shape pgm_vector = pgms.reshape(bsz * n_block, n_step, n_vocab) bsz, n_block, n_step, n_param = params.shape param_vector = params.reshape(bsz * n_block, n_step, n_param) index = (n_step - 1) * nd.ones(bsz * n_block).astype('int64') index = index.as_in_context(ctx) pred = executor(pgm_vector, param_vector, index) pred = nd.softmax(pred, axis=1) #print(pred.shape) pred = pred[:, 1] pred = pred.reshape(bsz, n_block, 32, 32, 32) rec = nd.max(pred, axis=1) #print("rec.shape:",rec.shape,"shapes.shape:",shapes.shape) #rec1 = rec.expand_dims(axis=1) rec1 = nd.log(rec + 1e-11) rec0 = nd.log(1 - rec + 1e-11) #rec_all = nd.concat(rec0, rec1, dim=1) #rec_all1 = nd.log(rec_all + 1e-10) #rec_all2 = nd.log(1-rec_all + 1e-10) gt = shapes.squeeze().astype('int64') loss = -nd.where(gt, rec1, rec0).mean(axis=(1, 2, 3)) #loss = -(nd.pick(rec_all1,gt,axis = 1,keepdims=True)).mean(axis = (1,2,3,4)) #loss = criterion(rec_all, gt) loss.backward() optimizer.step(loss.shape[0], ignore_stale_grad=True) l = loss.mean().asscalar() rendered_shapes = decode_multiple_block(pgms_int, params_int) rendered_shapes = nd.from_numpy(rendered_shapes).astype( 'float32').as_in_context(mx.cpu()) IoU2 = BatchIoU(raw_shapes, rendered_shapes) reconstruction = (rec.as_in_context(mx.cpu()) > 0.5).astype('float32') IoU1 = BatchIoU(reconstruction, raw_shapes) #print("IoU1:",IoU1,IoU2) IoU1 = IoU1.mean() IoU2 = IoU2.mean() end = time.time() if idx % opt.info_interval == 0: print( "Train: epoch {} batch {}/{}, loss = {:.3f}, IoU1 = {:.3f}, IoU2 = {:.3f}, time = {:.3f}" .format(epoch, idx, len(train_loader), l, IoU1, IoU2, end - start)) sys.stdout.flush()
def run(): # get options opt = options_guided_adaptation.parse() opt_gen = options_train_generator.parse() opt_exe = options_train_executor.parse() print('===== arguments: guided adaptation =====') for key, val in vars(opt).items(): print("{:20} {}".format(key, val)) print('===== arguments: guided adaptation =====') if not os.path.isdir(opt.save_folder): os.makedirs(opt.save_folder) # build loaders train_set = ShapeNet3D(opt.train_file) train_loader = gdata.DataLoader(dataset=train_set, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers) val_set = ShapeNet3D(opt.val_file) val_loader = gdata.DataLoader(dataset=val_set, batch_size=opt.batch_size, shuffle=False, num_workers=opt.num_workers) def visual(path, epoch, gen_shapes, file_name, nums_samples): data = gen_shapes.transpose((0, 3, 2, 1)) data = np.flip(data, axis=2) num_shapes = data.shape[0] for i in range(min(nums_samples, num_shapes)): voxels = data[i] save_name = os.path.join(path, file_name.format(epoch, i)) visualization(voxels, threshold=0.1, save_name=save_name, uniform_size=0.9) ctx = d2l.try_gpu() # load program generator generator = BlockOuterNet(opt_gen) generator.init_blocks(ctx) generator.load_parameters("model of blockouternet") # load program executor executor = RenderNet(opt_exe) executor.initialize(init=init.Xavier(), ctx=ctx) executor.load_parameters("model of executor") # build loss functions criterion = gloss.SoftmaxCrossEntropyLoss(axis=1, from_logits=True) optimizer = Trainer( generator.collect_params(), "adam", { "learning_rate": opt.learning_rate, "wd": opt.weight_decay, 'beta1': opt.beta1, 'beta2': opt.beta2, 'clip_gradient': opt.grad_clip }) print("###################") print("testing") gen_shapes, ori_shapes = validate(0, val_loader, generator, opt, ctx, gen_shape=True) #visual('imgs of chairs/adaption/chair/',0,ori_shapes,'GT {}-{}.png',8) #visual('imgs of chairs/adaption/chair/',0,gen_shapes,'epoch{}-{}.png',8) gen_shapes = nd.from_numpy(gen_shapes) ori_shapes = nd.from_numpy(ori_shapes) #print(gen_shapes.dtype,ori_shapes.dtype) #print("done",ori_shapes.shape,gen_shapes.shape) IoU = BatchIoU(gen_shapes, ori_shapes) #print(IoU) print("iou: ", IoU.mean()) best_iou = 0 print(opt.epochs) for epoch in range(1, opt.epochs + 1): print("###################") print("adaptation") train(epoch, train_loader, generator, executor, criterion, optimizer, opt, ctx) print("###################") print("testing") gen_shapes, ori_shapes = validate(epoch, val_loader, generator, opt, ctx, gen_shape=True) #visual('imgs of chairs/adaption/chair/',epoch,gen_shapes,'epoch{}-{}.png',8) gen_shapes = nd.from_numpy(gen_shapes) ori_shapes = nd.from_numpy(ori_shapes) IoU = BatchIoU(gen_shapes, ori_shapes) print("iou: ", IoU.mean()) if epoch % opt.save_interval == 0: print('Saving...') generator.save_parameters("generator of GA on shapenet") optimizer.save_states("optimazer of generator of GA on shapenet") if IoU.mean() >= best_iou: print('Saving best model') generator.save_parameters("generator of GA on shapenet") optimizer.save_states("optimazer of generator of GA on shapenet") best_iou = IoU.mean()
def train(epoch, train_loader, model, crit_cls, crit_reg, optimizer, opt, ctx): """ One epoch training """ cls_w = opt.cls_weight reg_w = opt.reg_weight # the prob: > 1 # the input of step t Operator where is missing FInferType attributeis always sampled from the output of step t-1 sample_prob = opt.inner_sample_prob for idx, data in enumerate(train_loader): start = time.time() #data, pgm, pgm_mask, param, param_mask shapes, labels, masks, params, param_masks = data[0], data[1], data[ 2], data[3], data[4] gt = shapes shapes = nd.expand_dims(shapes, axis=1) #print(labels[0],params[0]) shapes = shapes.as_in_context(ctx) labels = labels.as_in_context(ctx) labels2 = labels.as_in_context(ctx) masks = masks.as_in_context(ctx) params = params.as_in_context(ctx) param_masks = param_masks.as_in_context(ctx) #shapes.attach_grad(),labels.attach_grad() with autograd.record(): out = model(shapes, labels, sample_prob) #out = model.decode(shapes) # reshape bsz, n_block, n_step = labels.shape labels = labels.reshape(bsz, -1) masks = masks.reshape(bsz, -1) out_pgm = out[0].reshape(bsz, n_block * n_step, opt.program_size + 1) bsz, n_block, n_step, n_param = params.shape params = params.reshape(bsz, n_block * n_step, n_param) param_masks = param_masks.reshape(bsz, n_block * n_step, n_param) out_param = out[1].reshape(bsz, n_block * n_step, n_param) loss_cls, acc = crit_cls(out_pgm, labels, masks) loss_reg = crit_reg(out_param, params, param_masks) loss = cls_w * loss_cls + reg_w * loss_reg loss.backward() optimizer.step(bsz, ignore_stale_grad=True) loss_cls = loss_cls.mean().asscalar() loss_reg = loss_reg.mean().asscalar() end = time.time() if idx % (opt.info_interval * 10) == 0: out_1 = nd.round(out[0]).astype('int64') out_2 = nd.round(out[1]).astype('int64') pred = nd.from_numpy(decode_multiple_block( out_1, out_2)).astype("float32").as_in_context(mx.cpu()) IoU = BatchIoU(pred, gt) print( "Train: epoch {} batch {}/{},loss_cls = {:.3f},loss_reg = {:.3f},acc = {:.3f},IoU = {:.3f},time = {:.2f}" .format(epoch, idx, len(train_loader), loss_cls, loss_reg, acc[0].asscalar(), IoU.mean(), end - start)) sys.stdout.flush()