def evaluate_hand_draw_net(cfg): # Enable the inbuilt cudnn auto-tuner to find the best algorithm to use torch.backends.cudnn.benchmark = True IMG_SIZE = cfg.CONST.IMG_H, cfg.CONST.IMG_W CROP_SIZE = cfg.CONST.CROP_IMG_H, cfg.CONST.CROP_IMG_W eval_transforms = utils.data_transforms.Compose([ utils.data_transforms.CenterCrop(IMG_SIZE, CROP_SIZE), utils.data_transforms.RandomBackground(cfg.TEST.RANDOM_BG_COLOR_RANGE), utils.data_transforms.Normalize(mean=cfg.DATASET.MEAN, std=cfg.DATASET.STD), utils.data_transforms.ToTensor(), ]) # Set up networks encoder = Encoder(cfg) decoder = Decoder(cfg) azi_classes, ele_classes = int(360 / cfg.CONST.BIN_SIZE), int( 180 / cfg.CONST.BIN_SIZE) view_estimater = ViewEstimater(cfg, azi_classes=azi_classes, ele_classes=ele_classes) if torch.cuda.is_available(): encoder = torch.nn.DataParallel(encoder).cuda() decoder = torch.nn.DataParallel(decoder).cuda() view_estimater = torch.nn.DataParallel(view_estimater).cuda() # Load weight # Load weight for encoder, decoder print('[INFO] %s Loading reconstruction weights from %s ...' % (dt.now(), cfg.EVALUATE_HAND_DRAW.RECONSTRUCTION_WEIGHTS)) rec_checkpoint = torch.load(cfg.EVALUATE_HAND_DRAW.RECONSTRUCTION_WEIGHTS) encoder.load_state_dict(rec_checkpoint['encoder_state_dict']) decoder.load_state_dict(rec_checkpoint['decoder_state_dict']) print('[INFO] Best reconstruction result at epoch %d ...' % rec_checkpoint['epoch_idx']) # Load weight for view estimater print('[INFO] %s Loading view estimation weights from %s ...' % (dt.now(), cfg.EVALUATE_HAND_DRAW.VIEW_ESTIMATION_WEIGHTS)) view_checkpoint = torch.load( cfg.EVALUATE_HAND_DRAW.VIEW_ESTIMATION_WEIGHTS) view_estimater.load_state_dict( view_checkpoint['view_estimator_state_dict']) print('[INFO] Best view estimation result at epoch %d ...' % view_checkpoint['epoch_idx']) for img_path in os.listdir(cfg.EVALUATE_HAND_DRAW.INPUT_IMAGE_FOLDER): eval_id = int(img_path[:-4]) input_img_path = os.path.join( cfg.EVALUATE_HAND_DRAW.INPUT_IMAGE_FOLDER, img_path) print(input_img_path) evaluate_hand_draw_img(cfg, encoder, decoder, view_estimater, input_img_path, eval_transforms, eval_id)
class Visualization_demo(): def __init__(self, cfg, output_dir): self.encoder = Encoder(cfg) self.decoder = Decoder(cfg) self.refiner = Refiner(cfg) self.merger = Merger(cfg) checkpoint = torch.load(cfg.CHECKPOINT) encoder_state_dict = clean_state_dict(checkpoint['encoder_state_dict']) self.encoder.load_state_dict(encoder_state_dict) decoder_state_dict = clean_state_dict(checkpoint['decoder_state_dict']) self.decoder.load_state_dict(decoder_state_dict) if cfg.NETWORK.USE_REFINER: refiner_state_dict = clean_state_dict( checkpoint['refiner_state_dict']) self.refiner.load_state_dict(refiner_state_dict) if cfg.NETWORK.USE_MERGER: merger_state_dict = clean_state_dict( checkpoint['merger_state_dict']) self.merger.load_state_dict(merger_state_dict) if not os.path.exists(output_dir): os.makedirs(output_dir) self.output_dir = output_dir def run_on_images(self, imgs, sid, mid, iid, sampled_idx): dir1 = os.path.join(output_dir, str(sid), str(mid)) if not os.path.exists(dir1): os.makedirs(dir1) deprocess = imagenet_deprocess(rescale_image=False) image_features = self.encoder(imgs) raw_features, generated_volume = self.decoder(image_features) generated_volume = self.merger(raw_features, generated_volume) generated_volume = self.refiner(generated_volume) mesh = cubify(generated_volume, 0.3) # mesh = voxel_to_world(meshes) save_mesh = os.path.join(dir1, "%s_%s.obj" % (iid, sampled_idx)) verts, faces = mesh.get_mesh_verts_faces(0) save_obj(save_mesh, verts, faces) generated_volume = generated_volume.squeeze() img = image_to_numpy(deprocess(imgs[0][0])) save_img = os.path.join(dir1, "%02d.png" % (iid)) # cv2.imwrite(save_img, img[:, :, ::-1]) cv2.imwrite(save_img, img) img1 = image_to_numpy(deprocess(imgs[0][1])) save_img1 = os.path.join(dir1, "%02d.png" % (sampled_idx)) cv2.imwrite(save_img1, img1) # cv2.imwrite(save_img1, img1[:, :, ::-1]) get_volume_views(generated_volume, dir1, iid, sampled_idx)
def run(): """ Run the experiment. """ is_ptr = False np.random.seed(RANDOM_SEED) max_val, max_length, pairs = read_data(name="test") np.random.shuffle(pairs) training_pairs = [tensors_from_pair(pair) for pair in pairs] data_dim = max_val + 1 hidden_dim = embedding_dim = 256 encoder = Encoder(input_dim=data_dim, embedding_dim=embedding_dim, hidden_dim=hidden_dim).to(device) if is_ptr: decoder = PtrDecoder(output_dim=data_dim, embedding_dim=embedding_dim, hidden_dim=hidden_dim).to(device) else: decoder = AttnDecoder(output_dim=data_dim, embedding_dim=embedding_dim, hidden_dim=hidden_dim).to(device) checkpoint = load_checkpoint("ptr" if is_ptr else "vanilla") if checkpoint: encoder.load_state_dict(checkpoint["encoder"]) decoder.load_state_dict(checkpoint["decoder"]) else: print("Count not find checkpoint file.") permutation_count, nondecreasing_count = 0, 0 for i in range(len(training_pairs)): input_tensor, target_tensor = training_pairs[i] output_tensor = evaluate(encoder=encoder, decoder=decoder, input_tensor=training_pairs[i][0], is_ptr=is_ptr) target, output = list(np.asarray( input_tensor.data).squeeze()), output_tensor[:-1] if is_permutation(target, output): permutation_count += 1 if nondecreasing(output) == 0: nondecreasing_count += 1 print("Permutation: %s" % (permutation_count / len(training_pairs))) print("Nondecreasing: %s" % (nondecreasing_count / len(training_pairs)))
class Quantitative_analysis_demo(): def __init__(self, cfg, output_dir): self.encoder = Encoder(cfg) self.decoder = Decoder(cfg) self.refiner = Refiner(cfg) self.merger = Merger(cfg) # self.thresh = cfg.VOXEL_THRESH self.th = cfg.TEST.VOXEL_THRESH checkpoint = torch.load(cfg.CHECKPOINT) encoder_state_dict = clean_state_dict(checkpoint['encoder_state_dict']) self.encoder.load_state_dict(encoder_state_dict) decoder_state_dict = clean_state_dict(checkpoint['decoder_state_dict']) self.decoder.load_state_dict(decoder_state_dict) if cfg.NETWORK.USE_REFINER: refiner_state_dict = clean_state_dict( checkpoint['refiner_state_dict']) self.refiner.load_state_dict(refiner_state_dict) if cfg.NETWORK.USE_MERGER: merger_state_dict = clean_state_dict( checkpoint['merger_state_dict']) self.merger.load_state_dict(merger_state_dict) self.output_dir = output_dir def calculate_iou(self, imgs, GT_voxels, sid, mid, iid): dir1 = os.path.join(self.output_dir, str(sid), str(mid)) if not os.path.exists(dir1): os.makedirs(dir1) image_features = self.encoder(imgs) raw_features, generated_volume = self.decoder(image_features) generated_volume = self.merger(raw_features, generated_volume) generated_volume = self.refiner(generated_volume) generated_volume = generated_volume.squeeze() sample_iou = [] for th in self.th: _volume = torch.ge(generated_volume, th).float() intersection = torch.sum(_volume.mul(GT_voxels)).float() union = torch.sum(torch.ge(_volume.add(GT_voxels), 1)).float() sample_iou.append((intersection / union).item()) return sample_iou
def run(): """ Run the experiment. """ is_ptr = True np.random.seed(RANDOM_SEED) max_val, max_length, pairs = read_data(name="test") np.random.shuffle(pairs) training_pairs = [tensors_from_pair(pair) for pair in pairs] data_dim = max_val + 1 hidden_dim = embedding_dim = 256 encoder = Encoder(input_dim=data_dim, embedding_dim=embedding_dim, hidden_dim=hidden_dim).to(device) if is_ptr: decoder = PtrDecoder(output_dim=data_dim, embedding_dim=embedding_dim, hidden_dim=hidden_dim).to(device) else: decoder = AttnDecoder(output_dim=data_dim, embedding_dim=embedding_dim, hidden_dim=hidden_dim).to(device) checkpoint = load_checkpoint("ptr" if is_ptr else "vanilla") if checkpoint: encoder.load_state_dict(checkpoint["encoder"]) decoder.load_state_dict(checkpoint["decoder"]) else: print("Count not find checkpoint file.") for i in range(10): input_tensor, target_tensor = training_pairs[i] output_tensor = evaluate(encoder=encoder, decoder=decoder, input_tensor=training_pairs[i][0], is_ptr=is_ptr) print(list(np.asarray(input_tensor.data).squeeze()), output_tensor[:-1])
class System(catacomb.System): def __init__(self): hidden_dim = embedding_dim = 128 data_dim = 101 self.encoder = Encoder(input_dim=data_dim, embedding_dim=embedding_dim, hidden_dim=hidden_dim) self.decoder = PtrDecoder(output_dim=data_dim, embedding_dim=embedding_dim, hidden_dim=hidden_dim) checkpoint = torch.load('./e1i0.ckpt', map_location='cpu') self.encoder.load_state_dict(checkpoint['encoder']) self.decoder.load_state_dict(checkpoint['decoder']) def output(self, unsorted): unsorted_list = str_to_array(unsorted['list']) unsorted_tensor = tensor_from_list(unsorted_list) return str(evaluate(self.encoder, self.decoder, unsorted_tensor, True))
def test_single_img(cfg): encoder = Encoder(cfg) decoder = Decoder(cfg) refiner = Refiner(cfg) merger = Merger(cfg) cfg.CONST.WEIGHTS = 'D:/Pix2Vox/Pix2Vox/pretrained/Pix2Vox-A-ShapeNet.pth' checkpoint = torch.load(cfg.CONST.WEIGHTS, map_location=torch.device('cpu')) fix_checkpoint = {} fix_checkpoint['encoder_state_dict'] = OrderedDict((k.split('module.')[1:][0], v) for k, v in checkpoint['encoder_state_dict'].items()) fix_checkpoint['decoder_state_dict'] = OrderedDict((k.split('module.')[1:][0], v) for k, v in checkpoint['decoder_state_dict'].items()) fix_checkpoint['refiner_state_dict'] = OrderedDict((k.split('module.')[1:][0], v) for k, v in checkpoint['refiner_state_dict'].items()) fix_checkpoint['merger_state_dict'] = OrderedDict((k.split('module.')[1:][0], v) for k, v in checkpoint['merger_state_dict'].items()) epoch_idx = checkpoint['epoch_idx'] encoder.load_state_dict(fix_checkpoint['encoder_state_dict']) decoder.load_state_dict(fix_checkpoint['decoder_state_dict']) if cfg.NETWORK.USE_REFINER: print('Use refiner') refiner.load_state_dict(fix_checkpoint['refiner_state_dict']) if cfg.NETWORK.USE_MERGER: print('Use merger') merger.load_state_dict(fix_checkpoint['merger_state_dict']) encoder.eval() decoder.eval() refiner.eval() merger.eval() img1_path = 'D:/Pix2Vox/Pix2Vox/rand/minecraft.png' img1_np = cv2.imread(img1_path, cv2.IMREAD_UNCHANGED).astype(np.float32) / 255. sample = np.array([img1_np]) IMG_SIZE = cfg.CONST.IMG_H, cfg.CONST.IMG_W CROP_SIZE = cfg.CONST.CROP_IMG_H, cfg.CONST.CROP_IMG_W test_transforms = utils.data_transforms.Compose([ utils.data_transforms.CenterCrop(IMG_SIZE, CROP_SIZE), utils.data_transforms.RandomBackground(cfg.TEST.RANDOM_BG_COLOR_RANGE), utils.data_transforms.Normalize(mean=cfg.DATASET.MEAN, std=cfg.DATASET.STD), utils.data_transforms.ToTensor(), ]) rendering_images = test_transforms(rendering_images=sample) rendering_images = rendering_images.unsqueeze(0) with torch.no_grad(): image_features = encoder(rendering_images) raw_features, generated_volume = decoder(image_features) if cfg.NETWORK.USE_MERGER and epoch_idx >= cfg.TRAIN.EPOCH_START_USE_MERGER: generated_volume = merger(raw_features, generated_volume) else: generated_volume = torch.mean(generated_volume, dim=1) if cfg.NETWORK.USE_REFINER and epoch_idx >= cfg.TRAIN.EPOCH_START_USE_REFINER: generated_volume = refiner(generated_volume) generated_volume = generated_volume.squeeze(0) img_dir = 'D:/Pix2Vox/Pix2Vox/output' gv = generated_volume.cpu().numpy() gv_new = np.swapaxes(gv, 2, 1) print(gv_new) rendering_views = utils.binvox_visualization.get_volume_views(gv_new, os.path.join(img_dir), epoch_idx)
def main(config): print('Starting') checkpoints = config.checkpoint.parent.glob(config.checkpoint.name + '_*.pth') checkpoints = [c for c in checkpoints if extract_id(c) in config.decoders] assert len(checkpoints) >= 1, "No checkpoints found." model_config = torch.load(config.checkpoint.parent / 'args.pth')[0] encoder = Encoder(model_config.encoder) encoder.load_state_dict(torch.load(checkpoints[0])['encoder_state']) encoder.eval() encoder = encoder.cuda() generators = [] generator_ids = [] for checkpoint in checkpoints: decoder = Decoder(model_config.decoder) decoder.load_state_dict(torch.load(checkpoint)['decoder_state']) decoder.eval() decoder = decoder.cuda() generator = SampleGenerator(decoder, config.batch_size, wav_freq=config.rate) generators.append(generator) generator_ids.append(extract_id(checkpoint)) xs = [] assert config.out_dir is not None if len(config.sample_dir) == 1 and config.sample_dir[0].is_dir(): top = config.sample_dir[0] file_paths = list(top.glob('**/*.wav')) + list(top.glob('**/*.h5')) else: file_paths = config.sample_dir print("File paths to be used:", file_paths) for file_path in file_paths: if file_path.suffix == '.wav': data, rate = librosa.load(file_path, sr=config.rate) data = helper_functions.mu_law(data) elif file_path.suffix == '.h5': data = helper_functions.mu_law( h5py.File(file_path, 'r')['wav'][:] / (2**15)) if data.shape[-1] % config.rate != 0: data = data[:-(data.shape[-1] % config.rate)] assert data.shape[-1] % config.rate == 0 print(data.shape) else: raise Exception(f'Unsupported filetype {file_path}') if config.sample_len: data = data[:config.sample_len] else: config.sample_len = len(data) xs.append(torch.tensor(data).unsqueeze(0).float().cuda()) xs = torch.stack(xs).contiguous() print(f'xs size: {xs.size()}') def save(x, decoder_idx, filepath): wav = helper_functions.inv_mu_law(x.cpu().numpy()) print(f'X size: {x.shape}') print(f'X min: {x.min()}, max: {x.max()}') save_audio(wav.squeeze(), config.out_dir / str(decoder_idx) / filepath.with_suffix('.wav').name, rate=config.rate) yy = {} with torch.no_grad(): zz = [] for xs_batch in torch.split(xs, config.batch_size): zz += [encoder(xs_batch)] zz = torch.cat(zz, dim=0) for i, generator_id in enumerate(generator_ids): yy[generator_id] = [] generator = generators[i] for zz_batch in torch.split(zz, config.batch_size): print("Batch shape:", zz_batch.shape) splits = torch.split(zz_batch, config.split_size, -1) audio_data = [] generator.reset() for cond in tqdm.tqdm(splits): audio_data += [generator.generate(cond).cpu()] audio_data = torch.cat(audio_data, -1) yy[generator_id] += [audio_data] yy[generator_id] = torch.cat(yy[generator_id], dim=0) for sample_result, filepath in zip(yy[generator_id], file_paths): save(sample_result, generator_id, filepath) del generator
def test_net(cfg, epoch_idx=-1, output_dir=None, test_data_loader=None, \ test_writer=None, encoder=None, decoder=None, refiner=None, merger=None): # Enable the inbuilt cudnn auto-tuner to find the best algorithm to use torch.backends.cudnn.benchmark = True # Load taxonomies of dataset taxonomies = [] with open(cfg.DATASETS[cfg.DATASET.TEST_DATASET.upper()].TAXONOMY_FILE_PATH, encoding='utf-8') as file: taxonomies = json.loads(file.read()) taxonomies = {t['taxonomy_id']: t for t in taxonomies} # Set up data loader if test_data_loader is None: # Set up data augmentation IMG_SIZE = cfg.CONST.IMG_H, cfg.CONST.IMG_W CROP_SIZE = cfg.CONST.CROP_IMG_H, cfg.CONST.CROP_IMG_W test_transforms = utils.data_transforms.Compose([ utils.data_transforms.CenterCrop(IMG_SIZE, CROP_SIZE), utils.data_transforms.RandomBackground(cfg.TEST.RANDOM_BG_COLOR_RANGE), utils.data_transforms.Normalize(mean=cfg.DATASET.MEAN, std=cfg.DATASET.STD), utils.data_transforms.ToTensor(), ]) dataset_loader = utils.data_loaders.DATASET_LOADER_MAPPING[cfg.DATASET.TEST_DATASET](cfg) test_data_loader = torch.utils.data.DataLoader( dataset=dataset_loader.get_dataset(utils.data_loaders.DatasetType.TEST, cfg.CONST.N_VIEWS_RENDERING, test_transforms), batch_size=1, num_workers=1, pin_memory=True, shuffle=False) # Set up networks if decoder is None or encoder is None: encoder = Encoder(cfg) decoder = Decoder(cfg) refiner = Refiner(cfg) merger = Merger(cfg) if torch.cuda.is_available(): encoder = torch.nn.DataParallel(encoder).cuda() decoder = torch.nn.DataParallel(decoder).cuda() refiner = torch.nn.DataParallel(refiner).cuda() merger = torch.nn.DataParallel(merger).cuda() print('[INFO] %s Loading weights from %s ...' % (dt.now(), cfg.CONST.WEIGHTS)) checkpoint = torch.load(cfg.CONST.WEIGHTS) epoch_idx = checkpoint['epoch_idx'] encoder.load_state_dict(checkpoint['encoder_state_dict']) decoder.load_state_dict(checkpoint['decoder_state_dict']) if cfg.NETWORK.USE_REFINER: refiner.load_state_dict(checkpoint['refiner_state_dict']) if cfg.NETWORK.USE_MERGER: merger.load_state_dict(checkpoint['merger_state_dict']) # Set up loss functions bce_loss = torch.nn.BCELoss() # Testing loop n_samples = len(test_data_loader) test_iou = dict() encoder_losses = utils.network_utils.AverageMeter() refiner_losses = utils.network_utils.AverageMeter() # Switch models to evaluation mode encoder.eval() decoder.eval() refiner.eval() merger.eval() for sample_idx, (taxonomy_id, sample_name, rendering_images, ground_truth_volume) in enumerate(test_data_loader): taxonomy_id = taxonomy_id[0] if isinstance(taxonomy_id[0], str) else taxonomy_id[0].item() sample_name = sample_name[0] with torch.no_grad(): # Get data from data loader rendering_images = utils.network_utils.var_or_cuda(rendering_images) ground_truth_volume = utils.network_utils.var_or_cuda(ground_truth_volume) # Test the encoder, decoder, refiner and merger image_features = encoder(rendering_images) raw_features, generated_volume = decoder(image_features) if cfg.NETWORK.USE_MERGER and epoch_idx >= cfg.TRAIN.EPOCH_START_USE_MERGER: generated_volume = merger(raw_features, generated_volume) else: generated_volume = torch.mean(generated_volume, dim=1) encoder_loss = bce_loss(generated_volume, ground_truth_volume) * 10 if cfg.NETWORK.USE_REFINER and epoch_idx >= cfg.TRAIN.EPOCH_START_USE_REFINER: generated_volume = refiner(generated_volume) refiner_loss = bce_loss(generated_volume, ground_truth_volume) * 10 else: refiner_loss = encoder_loss print("vox shape {}".format(generated_volume.shape)) # Append loss and accuracy to average metrics encoder_losses.update(encoder_loss.item()) refiner_losses.update(refiner_loss.item()) # IoU per sample sample_iou = [] for th in cfg.TEST.VOXEL_THRESH: _volume = torch.ge(generated_volume, th).float() intersection = torch.sum(_volume.mul(ground_truth_volume)).float() union = torch.sum(torch.ge(_volume.add(ground_truth_volume), 1)).float() sample_iou.append((intersection / union).item()) # IoU per taxonomy if not taxonomy_id in test_iou: test_iou[taxonomy_id] = {'n_samples': 0, 'iou': []} test_iou[taxonomy_id]['n_samples'] += 1 test_iou[taxonomy_id]['iou'].append(sample_iou) # Append generated volumes to TensorBoard if output_dir and sample_idx < 3: img_dir = output_dir % 'images' # Volume Visualization gv = generated_volume.cpu().numpy() rendering_views = utils.binvox_visualization.get_volume_views(gv, os.path.join(img_dir, 'test'), epoch_idx) if not test_writer is None: test_writer.add_image('Test Sample#%02d/Volume Reconstructed' % sample_idx, rendering_views, epoch_idx) gtv = ground_truth_volume.cpu().numpy() rendering_views = utils.binvox_visualization.get_volume_views(gtv, os.path.join(img_dir, 'test'), epoch_idx) if not test_writer is None: test_writer.add_image('Test Sample#%02d/Volume GroundTruth' % sample_idx, rendering_views, epoch_idx) # Print sample loss and IoU print('[INFO] %s Test[%d/%d] Taxonomy = %s Sample = %s EDLoss = %.4f RLoss = %.4f IoU = %s' % \ (dt.now(), sample_idx + 1, n_samples, taxonomy_id, sample_name, encoder_loss.item(), \ refiner_loss.item(), ['%.4f' % si for si in sample_iou])) # Output testing results mean_iou = [] for taxonomy_id in test_iou: test_iou[taxonomy_id]['iou'] = np.mean(test_iou[taxonomy_id]['iou'], axis=0) mean_iou.append(test_iou[taxonomy_id]['iou'] * test_iou[taxonomy_id]['n_samples']) mean_iou = np.sum(mean_iou, axis=0) / n_samples # Print header print('============================ TEST RESULTS ============================') print('Taxonomy', end='\t') print('#Sample', end='\t') print('Baseline', end='\t') for th in cfg.TEST.VOXEL_THRESH: print('t=%.2f' % th, end='\t') print() # Print body for taxonomy_id in test_iou: print('%s' % taxonomies[taxonomy_id]['taxonomy_name'].ljust(8), end='\t') print('%d' % test_iou[taxonomy_id]['n_samples'], end='\t') if 'baseline' in taxonomies[taxonomy_id]: print('%.4f' % taxonomies[taxonomy_id]['baseline']['%d-view' % cfg.CONST.N_VIEWS_RENDERING], end='\t\t') else: print('N/a', end='\t\t') for ti in test_iou[taxonomy_id]['iou']: print('%.4f' % ti, end='\t') print() # Print mean IoU for each threshold print('Overall ', end='\t\t\t\t') for mi in mean_iou: print('%.4f' % mi, end='\t') print('\n') # Add testing results to TensorBoard max_iou = np.max(mean_iou) if not test_writer is None: test_writer.add_scalar('EncoderDecoder/EpochLoss', encoder_losses.avg, epoch_idx) test_writer.add_scalar('Refiner/EpochLoss', refiner_losses.avg, epoch_idx) test_writer.add_scalar('Refiner/IoU', max_iou, epoch_idx) return max_iou
def test_net(cfg, epoch_idx=-1, output_dir=None, test_data_loader=None, \ test_writer=None, encoder=None, decoder=None, refiner=None, merger=None): # Enable the inbuilt cudnn auto-tuner to find the best algorithm to use torch.backends.cudnn.benchmark = True # Load taxonomies of dataset taxonomies = [] with open( cfg.DATASETS[cfg.DATASET.TEST_DATASET.upper()].TAXONOMY_FILE_PATH, encoding='utf-8') as file: taxonomies = json.loads(file.read()) taxonomies = {t['taxonomy_id']: t for t in taxonomies} # Set up data loader if test_data_loader is None: # Set up data augmentation IMG_SIZE = cfg.CONST.IMG_H, cfg.CONST.IMG_W CROP_SIZE = cfg.CONST.CROP_IMG_H, cfg.CONST.CROP_IMG_W test_transforms = utils.data_transforms.Compose([ utils.data_transforms.CenterCrop(IMG_SIZE, CROP_SIZE), utils.data_transforms.RandomBackground( cfg.TEST.RANDOM_BG_COLOR_RANGE), utils.data_transforms.Normalize(mean=cfg.DATASET.MEAN, std=cfg.DATASET.STD), utils.data_transforms.ToTensor(), ]) dataset_loader = utils.data_loaders.DATASET_LOADER_MAPPING[ cfg.DATASET.TEST_DATASET](cfg) test_data_loader = torch.utils.data.DataLoader( dataset=dataset_loader.get_dataset( utils.data_loaders.DatasetType.TEST, cfg.CONST.N_VIEWS_RENDERING, test_transforms), batch_size=1, num_workers=1, pin_memory=True, shuffle=False) # Set up networks if decoder is None or encoder is None: encoder = Encoder(cfg) decoder = Decoder(cfg) refiner = Refiner(cfg) merger = Merger(cfg) if torch.cuda.is_available(): encoder = torch.nn.DataParallel(encoder).cuda() decoder = torch.nn.DataParallel(decoder).cuda() refiner = torch.nn.DataParallel(refiner).cuda() merger = torch.nn.DataParallel(merger).cuda() print('[INFO] %s Loading weights from %s ...' % (dt.now(), cfg.CONST.WEIGHTS)) if torch.cuda.is_available(): checkpoint = torch.load(cfg.CONST.WEIGHTS) else: map_location = torch.device('cpu') checkpoint = torch.load(cfg.CONST.WEIGHTS, map_location=map_location) epoch_idx = checkpoint['epoch_idx'] print('Epoch ID of the current model is {}'.format(epoch_idx)) encoder.load_state_dict(checkpoint['encoder_state_dict']) decoder.load_state_dict(checkpoint['decoder_state_dict']) if cfg.NETWORK.USE_REFINER: refiner.load_state_dict(checkpoint['refiner_state_dict']) if cfg.NETWORK.USE_MERGER: merger.load_state_dict(checkpoint['merger_state_dict']) # Set up loss functions bce_loss = torch.nn.BCELoss() # Testing loop n_samples = len(test_data_loader) test_iou = dict() encoder_losses = utils.network_utils.AverageMeter() refiner_losses = utils.network_utils.AverageMeter() # Switch models to evaluation mode encoder.eval() decoder.eval() refiner.eval() merger.eval() print("test data loader type is {}".format(type(test_data_loader))) for sample_idx, (taxonomy_id, sample_name, rendering_images) in enumerate(test_data_loader): taxonomy_id = taxonomy_id[0] if isinstance( taxonomy_id[0], str) else taxonomy_id[0].item() sample_name = sample_name[0] print("sample IDx {}".format(sample_idx)) print("taxonomy id {}".format(taxonomy_id)) with torch.no_grad(): # Get data from data loader rendering_images = utils.network_utils.var_or_cuda( rendering_images) print("Shape of the loaded images {}".format( rendering_images.shape)) # Test the encoder, decoder, refiner and merger image_features = encoder(rendering_images) raw_features, generated_volume = decoder(image_features) if cfg.NETWORK.USE_MERGER: generated_volume = merger(raw_features, generated_volume) else: generated_volume = torch.mean(generated_volume, dim=1) if cfg.NETWORK.USE_REFINER: generated_volume = refiner(generated_volume) print("vox shape {}".format(generated_volume.shape)) gv = generated_volume.cpu().numpy() rendering_views = utils.binvox_visualization.get_volume_views( gv, os.path.join('./LargeDatasets/inference_images/', 'inference'), sample_idx) print("gv shape is {}".format(gv.shape)) return gv, rendering_images
def main(): #parsing the arguments args, _ = parse_arguments() #setup logging #output_dir = Path('/content/drive/My Drive/image-captioning/output') output_dir = Path(args.output_directory) output_dir.mkdir(parents=True, exist_ok=True) logfile_path = Path(output_dir / "output.log") setup_logging(logfile=logfile_path) #setup and read config.ini #config_file = Path('/content/drive/My Drive/image-captioning/config.ini') config_file = Path('../config.ini') reading_config(config_file) #tensorboard tensorboard_logfile = Path(output_dir / 'tensorboard') tensorboard_writer = SummaryWriter(tensorboard_logfile) #load dataset #dataset_dir = Path('/content/drive/My Drive/Flickr8k_Dataset') dataset_dir = Path(args.dataset) images_path = Path(dataset_dir / Config.get("images_dir")) captions_path = Path(dataset_dir / Config.get("captions_dir")) training_loader, validation_loader, testing_loader = data_loaders( images_path, captions_path) #load the model (encoder, decoder, optimizer) embed_size = Config.get("encoder_embed_size") hidden_size = Config.get("decoder_hidden_size") batch_size = Config.get("training_batch_size") epochs = Config.get("epochs") feature_extraction = Config.get("feature_extraction") raw_captions = read_captions(captions_path) id_to_word, word_to_id = dictionary(raw_captions, threshold=5) vocab_size = len(id_to_word) encoder = Encoder(embed_size, feature_extraction) decoder = Decoder(embed_size, hidden_size, vocab_size, batch_size) #load pretrained embeddings #pretrained_emb_dir = Path('/content/drive/My Drive/word2vec') pretrained_emb_dir = Path(args.pretrained_embeddings) pretrained_emb_file = Path(pretrained_emb_dir / Config.get("pretrained_emb_path")) pretrained_embeddings = load_pretrained_embeddings(pretrained_emb_file, id_to_word) #load the optimizer learning_rate = Config.get("learning_rate") optimizer = adam_optimizer(encoder, decoder, learning_rate) #loss funciton criterion = cross_entropy #load checkpoint checkpoint_file = Path(output_dir / Config.get("checkpoint_file")) checkpoint_captioning = load_checkpoint(checkpoint_file) #using available device(gpu/cpu) encoder = encoder.to(Config.get("device")) decoder = decoder.to(Config.get("device")) pretrained_embeddings = pretrained_embeddings.to(Config.get("device")) start_epoch = 1 if checkpoint_captioning is not None: start_epoch = checkpoint_captioning['epoch'] + 1 encoder.load_state_dict(checkpoint_captioning['encoder']) decoder.load_state_dict(checkpoint_captioning['decoder']) optimizer.load_state_dict(checkpoint_captioning['optimizer']) logger.info( 'Initialized encoder, decoder and optimizer from loaded checkpoint' ) del checkpoint_captioning #image captioning model model = ImageCaptioning(encoder, decoder, optimizer, criterion, training_loader, validation_loader, testing_loader, pretrained_embeddings, output_dir, tensorboard_writer) #training and testing the model if args.training: validate_every = Config.get("validate_every") model.train(epochs, validate_every, start_epoch) elif args.testing: images_path = Path(images_path / Config.get("images_dir")) model.testing(id_to_word, images_path)
def test_img(cfg): encoder = Encoder(cfg) decoder = Decoder(cfg) refiner = Refiner(cfg) merger = Merger(cfg) cfg.CONST.WEIGHTS = '/Users/pranavpomalapally/Downloads/new-Pix2Vox-A-ShapeNet.pth' checkpoint = torch.load(cfg.CONST.WEIGHTS, map_location=torch.device('cpu')) print() # fix_checkpoint = {} # fix_checkpoint['encoder_state_dict'] = OrderedDict((k.split('module.')[1:][0], v) for k, v in checkpoint['encoder_state_dict'].items()) # fix_checkpoint['decoder_state_dict'] = OrderedDict((k.split('module.')[1:][0], v) for k, v in checkpoint['decoder_state_dict'].items()) # fix_checkpoint['refiner_state_dict'] = OrderedDict((k.split('module.')[1:][0], v) for k, v in checkpoint['refiner_state_dict'].items()) # fix_checkpoint['merger_state_dict'] = OrderedDict((k.split('module.')[1:][0], v) for k, v in checkpoint['merger_state_dict'].items()) # fix_checkpoint['encoder_state_dict'] = OrderedDict((k.split('module.')[0], v) for k, v in checkpoint['encoder_state_dict'].items()) # fix_checkpoint['decoder_state_dict'] = OrderedDict((k.split('module.')[0], v) for k, v in checkpoint['decoder_state_dict'].items()) # fix_checkpoint['refiner_state_dict'] = OrderedDict((k.split('module.')[0], v) for k, v in checkpoint['refiner_state_dict'].items()) # fix_checkpoint['merger_state_dict'] = OrderedDict((k.split('module.')[0], v) for k, v in checkpoint['merger_state_dict'].items()) epoch_idx = checkpoint['epoch_idx'] # encoder.load_state_dict(fix_checkpoint['encoder_state_dict']) # decoder.load_state_dict(fix_checkpoint['decoder_state_dict']) encoder.load_state_dict(checkpoint['encoder_state_dict']) decoder.load_state_dict(checkpoint['decoder_state_dict']) # if cfg.NETWORK.USE_REFINER: # print('Use refiner') # refiner.load_state_dict(fix_checkpoint['refiner_state_dict']) print('Use refiner') refiner.load_state_dict(checkpoint['refiner_state_dict']) if cfg.NETWORK.USE_MERGER: print('Use merger') # merger.load_state_dict(fix_checkpoint['merger_state_dict']) merger.load_state_dict(checkpoint['merger_state_dict']) encoder.eval() decoder.eval() refiner.eval() merger.eval() #img1_path = '/Users/pranavpomalapally/Downloads/ShapeNetRendering/02691156/1a04e3eab45ca15dd86060f189eb133/rendering/00.png' img1_path = '/Users/pranavpomalapally/Downloads/09 copy.png' img1_np = cv2.imread(img1_path, cv2.IMREAD_UNCHANGED).astype( np.float32) / 255. sample = np.array([img1_np]) IMG_SIZE = cfg.CONST.IMG_H, cfg.CONST.IMG_W CROP_SIZE = cfg.CONST.CROP_IMG_H, cfg.CONST.CROP_IMG_W test_transforms = utils.data_transforms.Compose([ utils.data_transforms.CenterCrop(IMG_SIZE, CROP_SIZE), utils.data_transforms.RandomBackground(cfg.TEST.RANDOM_BG_COLOR_RANGE), utils.data_transforms.Normalize(mean=cfg.DATASET.MEAN, std=cfg.DATASET.STD), utils.data_transforms.ToTensor(), ]) rendering_images = test_transforms(rendering_images=sample) rendering_images = rendering_images.unsqueeze(0) with torch.no_grad(): image_features = encoder(rendering_images) raw_features, generated_volume = decoder(image_features) if cfg.NETWORK.USE_MERGER and epoch_idx >= cfg.TRAIN.EPOCH_START_USE_MERGER: generated_volume = merger(raw_features, generated_volume) else: generated_volume = torch.mean(generated_volume, dim=1) # if cfg.NETWORK.USE_REFINER and epoch_idx >= cfg.TRAIN.EPOCH_START_USE_REFINER: # generated_volume = refiner(generated_volume) generated_volume = refiner(generated_volume) generated_volume = generated_volume.squeeze(0) img_dir = '/Users/pranavpomalapally/Downloads/outputs' # gv = generated_volume.cpu().numpy() gv = generated_volume.cpu().detach().numpy() gv_new = np.swapaxes(gv, 2, 1) os.environ['KMP_DUPLICATE_LIB_OK'] = 'True' rendering_views = utils.binvox_visualization.get_volume_views( gv_new, img_dir, epoch_idx)
class Trainer: def __init__(self, config): self.config = config self.config.data.n_datasets = len(config.data.datasets) print("No of datasets used:", self.config.data.n_datasets) torch.manual_seed(config.env.seed) torch.cuda.manual_seed(config.env.seed) self.expPath = self.config.env.expPath self.logger = Logger("Training", "logs/training.log") self.data = [ DatasetSet(data_path, config.data.seq_len, config.data) for data_path in config.data.datasets ] self.losses_recon = [ LossMeter(f'recon {i}') for i in range(self.config.data.n_datasets) ] self.loss_d_right = LossMeter('d') self.loss_total = [ LossMeter(f'total {i}') for i in range(self.config.data.n_datasets) ] self.evals_recon = [ LossMeter(f'recon {i}') for i in range(self.config.data.n_datasets) ] self.eval_d_right = LossMeter('eval d') self.eval_total = [ LossMeter(f'eval total {i}') for i in range(self.config.data.n_datasets) ] self.encoder = Encoder(config.encoder) self.decoders = torch.nn.ModuleList([ Decoder(config.decoder) for _ in range(self.config.data.n_datasets) ]) self.classifier = DomainClassifier( config.domain_classifier, num_classes=self.config.data.n_datasets) states = None if config.env.checkpoint: checkpoint_args_path = os.path.dirname( config.env.checkpoint) + '/args.pth' checkpoint_args = torch.load(checkpoint_args_path) self.start_epoch = checkpoint_args[-1] + 1 states = [ torch.load(self.config.env.checkpoint + f'_{i}.pth') for i in range(self.config.data.n_datasets) ] self.encoder.load_state_dict(states[0]['encoder_state']) for i in range(self.config.data.n_datasets): self.decoders[i].load_state_dict(states[i]['decoder_state']) self.classifier.load_state_dict(states[0]['discriminator_state']) self.logger.info('Loaded checkpoint parameters') raise NotImplementedError else: self.start_epoch = 0 self.encoder = torch.nn.DataParallel(self.encoder).cuda() self.classifier = torch.nn.DataParallel(self.classifier).cuda() for i, decoder in enumerate(self.decoders): self.decoders[i] = torch.nn.DataParallel(decoder).cuda() self.model_optimizers = [ optim.Adam(chain(self.encoder.parameters(), decoder.parameters()), lr=config.data.lr) for decoder in self.decoders ] self.classifier_optimizer = optim.Adam(self.classifier.parameters(), lr=config.data.lr) if config.env.checkpoint and config.env.load_optimizer: for i in range(self.config.data.n_datasets): self.model_optimizers[i].load_state_dict( states[i]['model_optimizer_state']) self.classifier_optimizer.load_state_dict( states[0]['d_optimizer_state']) self.lr_managers = [] for i in range(self.config.data.n_datasets): self.lr_managers.append( torch.optim.lr_scheduler.ExponentialLR( self.model_optimizers[i], config.data.lr_decay)) self.lr_managers[i].last_epoch = self.start_epoch self.lr_managers[i].step() def eval_batch(self, x, x_aug, dset_num): x, x_aug = x.float(), x_aug.float() z = self.encoder(x) y = self.decoders[dset_num](x, z) z_logits = self.classifier(z) z_classification = torch.max(z_logits, dim=1)[1] z_accuracy = (z_classification == dset_num).float().mean() self.eval_d_right.add(z_accuracy.data.item()) # discriminator_right = F.cross_entropy(z_logits, dset_num).mean() discriminator_right = F.cross_entropy( z_logits, torch.tensor([dset_num] * x.size(0)).long().cuda()).mean() recon_loss = cross_entropy_loss(y, x) self.evals_recon[dset_num].add(recon_loss.data.cpu().numpy().mean()) total_loss = discriminator_right.data.item() * self.config.domain_classifier.d_lambda + \ recon_loss.mean().data.item() self.eval_total[dset_num].add(total_loss) return total_loss def train_batch(self, x, x_aug, dset_num): x, x_aug = x.float(), x_aug.float() # Optimize D - classifier right z = self.encoder(x) z_logits = self.classifier(z) discriminator_right = F.cross_entropy( z_logits, torch.tensor([dset_num] * x.size(0)).long().cuda()).mean() loss = discriminator_right * self.config.domain_classifier.d_lambda self.loss_d_right.add(loss.data.item()) self.classifier_optimizer.zero_grad() loss.backward() if self.config.domain_classifier.grad_clip is not None: clip_grad_value_(self.classifier.parameters(), self.config.domain_classifier.grad_clip) self.classifier_optimizer.step() # optimize G - reconstructs well, classifier wrong z = self.encoder(x_aug) y = self.decoders[dset_num](x, z) z_logits = self.classifier(z) discriminator_wrong = -F.cross_entropy( z_logits, torch.tensor([dset_num] * x.size(0)).long().cuda()).mean() if not (-100 < discriminator_right.data.item() < 100): self.logger.debug(f'z_logits: {z_logits.detach().cpu().numpy()}') self.logger.debug(f'dset_num: {dset_num}') recon_loss = cross_entropy_loss(y, x) self.losses_recon[dset_num].add(recon_loss.data.cpu().numpy().mean()) loss = (recon_loss.mean() + self.config.domain_classifier.d_lambda * discriminator_wrong) self.model_optimizers[dset_num].zero_grad() loss.backward() if self.config.domain_classifier.grad_clip is not None: clip_grad_value_(self.encoder.parameters(), self.config.domain_classifier.grad_clip) clip_grad_value_(self.decoders[dset_num].parameters(), self.config.domain_classifier.grad_clip) self.model_optimizers[dset_num].step() self.loss_total[dset_num].add(loss.data.item()) return loss.data.item() def train_epoch(self, epoch): for meter in self.losses_recon: meter.reset() self.loss_d_right.reset() for i in range(len(self.loss_total)): self.loss_total[i].reset() self.encoder.train() self.classifier.train() for decoder in self.decoders: decoder.train() n_batches = self.config.data.epoch_len with tqdm(total=n_batches, desc='Train epoch %d' % epoch) as train_enum: for batch_num in range(n_batches): if self.config.data.short and batch_num == 3: break dset_num = batch_num % self.config.data.n_datasets x, x_aug = next(self.data[dset_num].train_iter) x = wrap_cuda(x) x_aug = wrap_cuda(x_aug) batch_loss = self.train_batch(x, x_aug, dset_num) train_enum.set_description( f'Train (loss: {batch_loss:.2f}) epoch {epoch}') train_enum.update() def evaluate_epoch(self, epoch): for meter in self.evals_recon: meter.reset() self.eval_d_right.reset() for i in range(len(self.eval_total)): self.eval_total[i].reset() self.encoder.eval() self.classifier.eval() for decoder in self.decoders: decoder.eval() n_batches = int(np.ceil(self.config.data.epoch_len / 10)) with tqdm(total=n_batches) as valid_enum, \ torch.no_grad(): for batch_num in range(n_batches): if self.config.data.short and batch_num == 10: break dset_num = batch_num % self.config.data.n_datasets x, x_aug = next(self.data[dset_num].valid_iter) x = wrap_cuda(x) x_aug = wrap_cuda(x_aug) batch_loss = self.eval_batch(x, x_aug, dset_num) valid_enum.set_description( f'Test (loss: {batch_loss:.2f}) epoch {epoch}') valid_enum.update() @staticmethod def format_losses(meters): losses = [meter.summarize_epoch() for meter in meters] return ', '.join('{:.4f}'.format(x) for x in losses) def train_losses(self): meters = [*self.losses_recon, self.loss_d_right] return self.format_losses(meters) def eval_losses(self): meters = [*self.evals_recon, self.eval_d_right] return self.format_losses(meters) def train(self): best_eval = [float('inf') for _ in range(self.config.data.n_datasets)] # Begin! for epoch in range(self.start_epoch, self.start_epoch + self.config.env.epochs): self.train_epoch(epoch) self.evaluate_epoch(epoch) self.logger.info(f'Epoch %s - Train loss: (%s), Test loss (%s)', epoch, self.train_losses(), self.eval_losses()) for i in range(len(self.lr_managers)): self.lr_managers[i].step() for dataset_id in range(self.config.data.n_datasets): val_loss = self.eval_total[dataset_id].summarize_epoch() if val_loss < best_eval[dataset_id]: self.save_model(f'bestmodel_{dataset_id}.pth', dataset_id) best_eval[dataset_id] = val_loss if not self.config.env.save_per_epoch: self.save_model(f'lastmodel_{dataset_id}.pth', dataset_id) else: self.save_model(f'lastmodel_{epoch}_rank_{dataset_id}.pth', dataset_id) torch.save([self.config, epoch], '%s/args.pth' % self.expPath) self.logger.debug('Ended epoch') def save_model(self, filename, decoder_id): save_path = self.expPath / filename torch.save( { 'encoder_state': self.encoder.module.state_dict(), 'decoder_state': self.decoders[decoder_id].module.state_dict(), 'discriminator_state': self.classifier.module.state_dict(), 'model_optimizer_state': self.model_optimizers[decoder_id].state_dict(), 'dataset': decoder_id, 'd_optimizer_state': self.classifier_optimizer.state_dict() }, save_path) self.logger.debug(f'Saved model to {save_path}')
def train_net(cfg): # Enable the inbuilt cudnn auto-tuner to find the best algorithm to use torch.backends.cudnn.benchmark = True # Set up data augmentation IMG_SIZE = cfg.CONST.IMG_H, cfg.CONST.IMG_W CROP_SIZE = cfg.CONST.CROP_IMG_H, cfg.CONST.CROP_IMG_W train_transforms = utils.data_transforms.Compose([ utils.data_transforms.RandomCrop(IMG_SIZE, CROP_SIZE), utils.data_transforms.RandomBackground( cfg.TRAIN.RANDOM_BG_COLOR_RANGE), utils.data_transforms.ColorJitter(cfg.TRAIN.BRIGHTNESS, cfg.TRAIN.CONTRAST, cfg.TRAIN.SATURATION), utils.data_transforms.RandomNoise(cfg.TRAIN.NOISE_STD), utils.data_transforms.Normalize(mean=cfg.DATASET.MEAN, std=cfg.DATASET.STD), utils.data_transforms.RandomFlip(), utils.data_transforms.RandomPermuteRGB(), utils.data_transforms.ToTensor(), ]) val_transforms = utils.data_transforms.Compose([ utils.data_transforms.CenterCrop(IMG_SIZE, CROP_SIZE), utils.data_transforms.RandomBackground(cfg.TEST.RANDOM_BG_COLOR_RANGE), utils.data_transforms.Normalize(mean=cfg.DATASET.MEAN, std=cfg.DATASET.STD), utils.data_transforms.ToTensor(), ]) # Set up data loader train_dataset_loader = utils.data_loaders.DATASET_LOADER_MAPPING[ cfg.DATASET.TRAIN_DATASET](cfg) val_dataset_loader = utils.data_loaders.DATASET_LOADER_MAPPING[ cfg.DATASET.TEST_DATASET](cfg) train_data_loader = torch.utils.data.DataLoader( dataset=train_dataset_loader.get_dataset( utils.data_loaders.DatasetType.TRAIN, cfg.CONST.N_VIEWS_RENDERING, train_transforms), batch_size=cfg.CONST.BATCH_SIZE, num_workers=cfg.TRAIN.NUM_WORKER, pin_memory=True, shuffle=True, drop_last=True) val_data_loader = torch.utils.data.DataLoader( dataset=val_dataset_loader.get_dataset( utils.data_loaders.DatasetType.VAL, cfg.CONST.N_VIEWS_RENDERING, val_transforms), batch_size=1, num_workers=1, pin_memory=True, shuffle=False) # Set up networks encoder = Encoder(cfg) decoder = Decoder(cfg) refiner = Refiner(cfg) merger = Merger(cfg) print('[DEBUG] %s Parameters in Encoder: %d.' % (dt.now(), utils.network_utils.count_parameters(encoder))) print('[DEBUG] %s Parameters in Decoder: %d.' % (dt.now(), utils.network_utils.count_parameters(decoder))) print('[DEBUG] %s Parameters in Refiner: %d.' % (dt.now(), utils.network_utils.count_parameters(refiner))) print('[DEBUG] %s Parameters in Merger: %d.' % (dt.now(), utils.network_utils.count_parameters(merger))) # Initialize weights of networks encoder.apply(utils.network_utils.init_weights) decoder.apply(utils.network_utils.init_weights) refiner.apply(utils.network_utils.init_weights) merger.apply(utils.network_utils.init_weights) # Set up solver if cfg.TRAIN.POLICY == 'adam': encoder_solver = torch.optim.Adam(filter(lambda p: p.requires_grad, encoder.parameters()), lr=cfg.TRAIN.ENCODER_LEARNING_RATE, betas=cfg.TRAIN.BETAS) decoder_solver = torch.optim.Adam(decoder.parameters(), lr=cfg.TRAIN.DECODER_LEARNING_RATE, betas=cfg.TRAIN.BETAS) refiner_solver = torch.optim.Adam(refiner.parameters(), lr=cfg.TRAIN.REFINER_LEARNING_RATE, betas=cfg.TRAIN.BETAS) merger_solver = torch.optim.Adam(merger.parameters(), lr=cfg.TRAIN.MERGER_LEARNING_RATE, betas=cfg.TRAIN.BETAS) elif cfg.TRAIN.POLICY == 'sgd': encoder_solver = torch.optim.SGD(filter(lambda p: p.requires_grad, encoder.parameters()), lr=cfg.TRAIN.ENCODER_LEARNING_RATE, momentum=cfg.TRAIN.MOMENTUM) decoder_solver = torch.optim.SGD(decoder.parameters(), lr=cfg.TRAIN.DECODER_LEARNING_RATE, momentum=cfg.TRAIN.MOMENTUM) refiner_solver = torch.optim.SGD(refiner.parameters(), lr=cfg.TRAIN.REFINER_LEARNING_RATE, momentum=cfg.TRAIN.MOMENTUM) merger_solver = torch.optim.SGD(merger.parameters(), lr=cfg.TRAIN.MERGER_LEARNING_RATE, momentum=cfg.TRAIN.MOMENTUM) else: raise Exception('[FATAL] %s Unknown optimizer %s.' % (dt.now(), cfg.TRAIN.POLICY)) # Set up learning rate scheduler to decay learning rates dynamically encoder_lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( encoder_solver, milestones=cfg.TRAIN.ENCODER_LR_MILESTONES, gamma=cfg.TRAIN.GAMMA) decoder_lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( decoder_solver, milestones=cfg.TRAIN.DECODER_LR_MILESTONES, gamma=cfg.TRAIN.GAMMA) refiner_lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( refiner_solver, milestones=cfg.TRAIN.REFINER_LR_MILESTONES, gamma=cfg.TRAIN.GAMMA) merger_lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( merger_solver, milestones=cfg.TRAIN.MERGER_LR_MILESTONES, gamma=cfg.TRAIN.GAMMA) if torch.cuda.is_available(): encoder = torch.nn.DataParallel(encoder).cuda() decoder = torch.nn.DataParallel(decoder).cuda() refiner = torch.nn.DataParallel(refiner).cuda() merger = torch.nn.DataParallel(merger).cuda() # Set up loss functions bce_loss = torch.nn.BCELoss() # Load pretrained model if exists init_epoch = 0 best_iou = -1 best_epoch = -1 if 'WEIGHTS' in cfg.CONST and cfg.TRAIN.RESUME_TRAIN: print('[INFO] %s Recovering from %s ...' % (dt.now(), cfg.CONST.WEIGHTS)) checkpoint = torch.load(cfg.CONST.WEIGHTS) init_epoch = checkpoint['epoch_idx'] best_iou = checkpoint['best_iou'] best_epoch = checkpoint['best_epoch'] encoder.load_state_dict(checkpoint['encoder_state_dict']) decoder.load_state_dict(checkpoint['decoder_state_dict']) if cfg.NETWORK.USE_REFINER: refiner.load_state_dict(checkpoint['refiner_state_dict']) if cfg.NETWORK.USE_MERGER: merger.load_state_dict(checkpoint['merger_state_dict']) print('[INFO] %s Recover complete. Current epoch #%d, Best IoU = %.4f at epoch #%d.' \ % (dt.now(), init_epoch, best_iou, best_epoch)) # Summary writer for TensorBoard output_dir = os.path.join(cfg.DIR.OUT_PATH, '%s', dt.now().isoformat()) log_dir = output_dir % 'logs' ckpt_dir = output_dir % 'checkpoints' train_writer = SummaryWriter(os.path.join(log_dir, 'train')) val_writer = SummaryWriter(os.path.join(log_dir, 'test')) # Training loop for epoch_idx in range(init_epoch, cfg.TRAIN.NUM_EPOCHES): # Tick / tock epoch_start_time = time() # Batch average meterics batch_time = utils.network_utils.AverageMeter() data_time = utils.network_utils.AverageMeter() encoder_losses = utils.network_utils.AverageMeter() refiner_losses = utils.network_utils.AverageMeter() # Adjust learning rate encoder_lr_scheduler.step() decoder_lr_scheduler.step() refiner_lr_scheduler.step() merger_lr_scheduler.step() # switch models to training mode encoder.train() decoder.train() merger.train() refiner.train() batch_end_time = time() n_batches = len(train_data_loader) for batch_idx, (taxonomy_names, sample_names, rendering_images, ground_truth_volumes) in enumerate(train_data_loader): # Measure data time data_time.update(time() - batch_end_time) # Get data from data loader rendering_images = utils.network_utils.var_or_cuda( rendering_images) ground_truth_volumes = utils.network_utils.var_or_cuda( ground_truth_volumes) # Train the encoder, decoder, refiner, and merger image_features = encoder(rendering_images) raw_features, generated_volumes = decoder(image_features) if cfg.NETWORK.USE_MERGER and epoch_idx >= cfg.TRAIN.EPOCH_START_USE_MERGER: generated_volumes = merger(raw_features, generated_volumes) else: generated_volumes = torch.mean(generated_volumes, dim=1) encoder_loss = bce_loss(generated_volumes, ground_truth_volumes) * 10 if cfg.NETWORK.USE_REFINER and epoch_idx >= cfg.TRAIN.EPOCH_START_USE_REFINER: generated_volumes = refiner(generated_volumes) refiner_loss = bce_loss(generated_volumes, ground_truth_volumes) * 10 else: refiner_loss = encoder_loss # Gradient decent encoder.zero_grad() decoder.zero_grad() refiner.zero_grad() merger.zero_grad() if cfg.NETWORK.USE_REFINER and epoch_idx >= cfg.TRAIN.EPOCH_START_USE_REFINER: encoder_loss.backward(retain_graph=True) refiner_loss.backward() else: encoder_loss.backward() encoder_solver.step() decoder_solver.step() refiner_solver.step() merger_solver.step() # Append loss to average metrics encoder_losses.update(encoder_loss.item()) refiner_losses.update(refiner_loss.item()) # Append loss to TensorBoard n_itr = epoch_idx * n_batches + batch_idx train_writer.add_scalar('EncoderDecoder/BatchLoss', encoder_loss.item(), n_itr) train_writer.add_scalar('Refiner/BatchLoss', refiner_loss.item(), n_itr) # Tick / tock batch_time.update(time() - batch_end_time) batch_end_time = time() print('[INFO] %s [Epoch %d/%d][Batch %d/%d] BatchTime = %.3f (s) DataTime = %.3f (s) EDLoss = %.4f RLoss = %.4f' % \ (dt.now(), epoch_idx + 1, cfg.TRAIN.NUM_EPOCHES, batch_idx + 1, n_batches, \ batch_time.val, data_time.val, encoder_loss.item(), refiner_loss.item())) # Append epoch loss to TensorBoard train_writer.add_scalar('EncoderDecoder/EpochLoss', encoder_losses.avg, epoch_idx + 1) train_writer.add_scalar('Refiner/EpochLoss', refiner_losses.avg, epoch_idx + 1) # Tick / tock epoch_end_time = time() print('[INFO] %s Epoch [%d/%d] EpochTime = %.3f (s) EDLoss = %.4f RLoss = %.4f' % (dt.now(), epoch_idx + 1, cfg.TRAIN.NUM_EPOCHES, epoch_end_time - epoch_start_time, \ encoder_losses.avg, refiner_losses.avg)) # Update Rendering Views if cfg.TRAIN.UPDATE_N_VIEWS_RENDERING: n_views_rendering = random.randint(1, cfg.CONST.N_VIEWS_RENDERING) train_data_loader.dataset.set_n_views_rendering(n_views_rendering) print('[INFO] %s Epoch [%d/%d] Update #RenderingViews to %d' % \ (dt.now(), epoch_idx + 2, cfg.TRAIN.NUM_EPOCHES, n_views_rendering)) # Validate the training models iou = test_net(cfg, epoch_idx + 1, output_dir, val_data_loader, val_writer, encoder, decoder, refiner, merger) # Save weights to file if (epoch_idx + 1) % cfg.TRAIN.SAVE_FREQ == 0: if not os.path.exists(ckpt_dir): os.makedirs(ckpt_dir) utils.network_utils.save_checkpoints(cfg, \ os.path.join(ckpt_dir, 'ckpt-epoch-%04d.pth' % (epoch_idx + 1)), \ epoch_idx + 1, encoder, encoder_solver, decoder, decoder_solver, \ refiner, refiner_solver, merger, merger_solver, best_iou, best_epoch) if iou > best_iou: if not os.path.exists(ckpt_dir): os.makedirs(ckpt_dir) best_iou = iou best_epoch = epoch_idx + 1 utils.network_utils.save_checkpoints(cfg, \ os.path.join(ckpt_dir, 'best-ckpt.pth'), \ epoch_idx + 1, encoder, encoder_solver, decoder, decoder_solver, \ refiner, refiner_solver, merger, merger_solver, best_iou, best_epoch) # Close SummaryWriter for TensorBoard train_writer.close() val_writer.close()
def test_single_img_net(cfg): encoder = Encoder(cfg) decoder = Decoder(cfg) refiner = Refiner(cfg) merger = Merger(cfg) print('[INFO] %s Loading weights from %s ...' % (dt.now(), cfg.CONST.WEIGHTS)) checkpoint = torch.load(cfg.CONST.WEIGHTS, map_location=torch.device('cpu')) fix_checkpoint = {} fix_checkpoint['encoder_state_dict'] = OrderedDict( (k.split('module.')[1:][0], v) for k, v in checkpoint['encoder_state_dict'].items()) fix_checkpoint['decoder_state_dict'] = OrderedDict( (k.split('module.')[1:][0], v) for k, v in checkpoint['decoder_state_dict'].items()) fix_checkpoint['refiner_state_dict'] = OrderedDict( (k.split('module.')[1:][0], v) for k, v in checkpoint['refiner_state_dict'].items()) fix_checkpoint['merger_state_dict'] = OrderedDict( (k.split('module.')[1:][0], v) for k, v in checkpoint['merger_state_dict'].items()) epoch_idx = checkpoint['epoch_idx'] encoder.load_state_dict(fix_checkpoint['encoder_state_dict']) decoder.load_state_dict(fix_checkpoint['decoder_state_dict']) if cfg.NETWORK.USE_REFINER: print('Use refiner') refiner.load_state_dict(fix_checkpoint['refiner_state_dict']) if cfg.NETWORK.USE_MERGER: print('Use merger') merger.load_state_dict(fix_checkpoint['merger_state_dict']) encoder.eval() decoder.eval() refiner.eval() merger.eval() img1_path = '/media/caig/FECA2C89CA2C406F/dataset/ShapeNetRendering_copy/03001627/1a74a83fa6d24b3cacd67ce2c72c02e/rendering/00.png' img1_np = cv2.imread(img1_path, cv2.IMREAD_UNCHANGED).astype( np.float32) / 255. sample = np.array([img1_np]) IMG_SIZE = cfg.CONST.IMG_H, cfg.CONST.IMG_W CROP_SIZE = cfg.CONST.CROP_IMG_H, cfg.CONST.CROP_IMG_W test_transforms = utils.data_transforms.Compose([ utils.data_transforms.CenterCrop(IMG_SIZE, CROP_SIZE), utils.data_transforms.RandomBackground(cfg.TEST.RANDOM_BG_COLOR_RANGE), utils.data_transforms.Normalize(mean=cfg.DATASET.MEAN, std=cfg.DATASET.STD), utils.data_transforms.ToTensor(), ]) rendering_images = test_transforms(rendering_images=sample) rendering_images = rendering_images.unsqueeze(0) with torch.no_grad(): image_features = encoder(rendering_images) raw_features, generated_volume = decoder(image_features) if cfg.NETWORK.USE_MERGER and epoch_idx >= cfg.TRAIN.EPOCH_START_USE_MERGER: generated_volume = merger(raw_features, generated_volume) else: generated_volume = torch.mean(generated_volume, dim=1) if cfg.NETWORK.USE_REFINER and epoch_idx >= cfg.TRAIN.EPOCH_START_USE_REFINER: generated_volume = refiner(generated_volume) generated_volume = generated_volume.squeeze(0) img_dir = '/media/caig/FECA2C89CA2C406F/sketch3D/sketch3D/test_output' gv = generated_volume.cpu().numpy() gv_new = np.swapaxes(gv, 2, 1) rendering_views = utils.binvox_visualization.get_volume_views( gv_new, os.path.join(img_dir), epoch_idx)
vocab = Vocabulary("./captions.json", args.NUM_CAPTIONS, num_fre=args.NUM_FRE) VOCAB_SIZE = vocab.num_words SEQ_LEN = vocab.max_sentence_len encoder = Encoder(args.ENCODER_OUTPUT_SIZE) decoder = Decoder(embed_size=args.EMBED_SIZE, hidden_size=args.HIDDEN_SIZE, attention_size=args.ATTENTION_SIZE, vocab_size=VOCAB_SIZE, encoder_size=2048, device=device, seq_len=SEQ_LEN + 2) encoder.load_state_dict(torch.load(args.ENCODER_MODEL_LOAD_PATH)) decoder.load_state_dict(torch.load(args.DECODER_MODEL_LOAD_PATH)) encoder.to(device) decoder.to(device) encoder.eval() decoder.eval() result_json = {"images": []} for path in IMG_PATH: img_name = path.split("/")[-1] img = Image.open(path) img = transform(img).unsqueeze(0).to( device) # [BATCH_SIZE(1) * CHANNEL * INPUT_SIZE * INPUT_SIZE] num_sentence = args.NUM_TOP_PROB top_prev_prob = torch.zeros((num_sentence, 1)).to(device)
def test_net(cfg, model_type, dataset_type, results_file_name, epoch_idx=-1, test_data_loader=None, test_writer=None, encoder=None, decoder=None, refiner=None, merger=None, save_results_to_file=False, show_voxels=False, path_to_times_csv=None): if model_type == Pix2VoxTypes.Pix2Vox_A or model_type == Pix2VoxTypes.Pix2Vox_Plus_Plus_A: use_refiner = True else: use_refiner = False # Enable the inbuilt cudnn auto-tuner to find the best algorithm to use torch.backends.cudnn.benchmark = True # Set up data loader if test_data_loader is None: # Set up data augmentation IMG_SIZE = cfg.CONST.IMG_H, cfg.CONST.IMG_W CROP_SIZE = cfg.CONST.CROP_IMG_H, cfg.CONST.CROP_IMG_W test_transforms = utils.data_transforms.Compose([ utils.data_transforms.CenterCrop(IMG_SIZE, CROP_SIZE), utils.data_transforms.RandomBackground( cfg.TEST.RANDOM_BG_COLOR_RANGE), utils.data_transforms.Normalize(mean=cfg.DATASET.MEAN, std=cfg.DATASET.STD), utils.data_transforms.ToTensor(), ]) dataset_loader = utils.data_loaders.DATASET_LOADER_MAPPING[ cfg.DATASET.TEST_DATASET](cfg) test_data_loader = torch.utils.data.DataLoader( dataset=dataset_loader.get_dataset(dataset_type, cfg.CONST.N_VIEWS_RENDERING, test_transforms), batch_size=1, num_workers=cfg.CONST.NUM_WORKER, pin_memory=True, shuffle=False) # Set up networks if decoder is None or encoder is None: encoder = Encoder(cfg, model_type) decoder = Decoder(cfg, model_type) if use_refiner: refiner = Refiner(cfg) merger = Merger(cfg, model_type) if torch.cuda.is_available(): encoder = torch.nn.DataParallel(encoder).cuda() decoder = torch.nn.DataParallel(decoder).cuda() if use_refiner: refiner = torch.nn.DataParallel(refiner).cuda() merger = torch.nn.DataParallel(merger).cuda() logging.info('Loading weights from %s ...' % (cfg.CONST.WEIGHTS)) checkpoint = torch.load(cfg.CONST.WEIGHTS) epoch_idx = checkpoint['epoch_idx'] encoder.load_state_dict(checkpoint['encoder_state_dict']) decoder.load_state_dict(checkpoint['decoder_state_dict']) if use_refiner: refiner.load_state_dict(checkpoint['refiner_state_dict']) if cfg.NETWORK.USE_MERGER: merger.load_state_dict(checkpoint['merger_state_dict']) # Set up loss functions bce_loss = torch.nn.BCELoss() # Testing loop n_samples = len(test_data_loader) test_iou = dict() encoder_losses = AverageMeter() if use_refiner: refiner_losses = AverageMeter() # Switch models to evaluation mode encoder.eval() decoder.eval() if use_refiner: refiner.eval() merger.eval() samples_names = [] edlosses = [] rlosses = [] ious_dict = {} for iou_threshold in cfg.TEST.VOXEL_THRESH: ious_dict[iou_threshold] = [] if path_to_times_csv is not None: n_view_list = [] times_list = [] for sample_idx, (taxonomy_id, sample_name, rendering_images, ground_truth_volume) in enumerate(test_data_loader): taxonomy_id = taxonomy_id[0] if isinstance( taxonomy_id[0], str) else taxonomy_id[0].item() sample_name = sample_name[0] with torch.no_grad(): # Get data from data loader rendering_images = utils.helpers.var_or_cuda(rendering_images) ground_truth_volume = utils.helpers.var_or_cuda( ground_truth_volume) if path_to_times_csv is not None: start_time = time.time() # Test the encoder, decoder, refiner and merger image_features = encoder(rendering_images) raw_features, generated_volume = decoder(image_features) if cfg.NETWORK.USE_MERGER and epoch_idx >= cfg.TRAIN.EPOCH_START_USE_MERGER: generated_volume = merger(raw_features, generated_volume) else: generated_volume = torch.mean(generated_volume, dim=1) encoder_loss = bce_loss(generated_volume, ground_truth_volume) * 10 if use_refiner and epoch_idx >= cfg.TRAIN.EPOCH_START_USE_REFINER: generated_volume = refiner(generated_volume) refiner_loss = bce_loss(generated_volume, ground_truth_volume) * 10 else: refiner_loss = encoder_loss if path_to_times_csv is not None: end_time = time.time() n_view_list.append(rendering_images.size()[1]) times_list.append(end_time - start_time) # Append loss and accuracy to average metrics encoder_losses.update(encoder_loss.item()) if use_refiner: refiner_losses.update(refiner_loss.item()) # IoU per sample sample_iou = [] for th in cfg.TEST.VOXEL_THRESH: _volume = torch.ge(generated_volume, th).float() intersection = torch.sum( _volume.mul(ground_truth_volume)).float() union = torch.sum(torch.ge(_volume.add(ground_truth_volume), 1)).float() sample_iou.append((intersection / union).item()) ious_dict[th].append((intersection / union).item()) # IoU per taxonomy if taxonomy_id not in test_iou: test_iou[taxonomy_id] = {'n_samples': 0, 'iou': []} test_iou[taxonomy_id]['n_samples'] += 1 test_iou[taxonomy_id]['iou'].append(sample_iou) # Append generated volumes to TensorBoard if show_voxels: with open("model.binvox", "wb") as f: v = br.Voxels( torch.ge(generated_volume, 0.2).float().cpu().numpy()[0], (32, 32, 32), (0, 0, 0), 1, "xyz") v.write(f) subprocess.run([VIEWVOX_EXE, "model.binvox"]) with open("model.binvox", "wb") as f: v = br.Voxels(ground_truth_volume.cpu().numpy()[0], (32, 32, 32), (0, 0, 0), 1, "xyz") v.write(f) subprocess.run([VIEWVOX_EXE, "model.binvox"]) # Print sample loss and IoU logging.info( 'Test[%d/%d] Taxonomy = %s Sample = %s EDLoss = %.4f RLoss = %.4f IoU = %s' % (sample_idx + 1, n_samples, taxonomy_id, sample_name, encoder_loss.item(), refiner_loss.item(), ['%.4f' % si for si in sample_iou])) samples_names.append(sample_name) edlosses.append(encoder_loss.item()) if use_refiner: rlosses.append(refiner_loss.item()) if save_results_to_file: save_test_results_to_csv(samples_names, edlosses, rlosses, ious_dict, path_to_csv=results_file_name) if path_to_times_csv is not None: save_times_to_csv(times_list, n_view_list, path_to_csv=path_to_times_csv) # Output testing results mean_iou = [] for taxonomy_id in test_iou: test_iou[taxonomy_id]['iou'] = np.mean(test_iou[taxonomy_id]['iou'], axis=0) mean_iou.append(test_iou[taxonomy_id]['iou'] * test_iou[taxonomy_id]['n_samples']) mean_iou = np.sum(mean_iou, axis=0) / n_samples # Print header print( '============================ TEST RESULTS ============================' ) print('Taxonomy', end='\t') print('#Sample', end='\t') print('Baseline', end='\t') for th in cfg.TEST.VOXEL_THRESH: print('t=%.2f' % th, end='\t') print() # Print mean IoU for each threshold print('Overall ', end='\t\t\t\t') for mi in mean_iou: print('%.4f' % mi, end='\t') print('\n') # Add testing results to TensorBoard max_iou = np.max(mean_iou) if test_writer is not None: test_writer.add_scalar('EncoderDecoder/EpochLoss', encoder_losses.avg, epoch_idx) if use_refiner: test_writer.add_scalar('Refiner/EpochLoss', refiner_losses.avg, epoch_idx) test_writer.add_scalar('Refiner/IoU', max_iou, epoch_idx) return max_iou