def main(dataset, image_path, model_path, cuda, crf): CONFIG = { 'voc12': { 'path_pytorch_model': 'data/models/deeplab_resnet101/voc12/deeplabv2_resnet101_VOC2012.pth', 'label_list': 'data/datasets/voc12/labels.txt', 'n_classes': 21, 'n_blocks': [3, 4, 23, 3], 'pyramids': [6, 3, 2, 1], 'image': { 'size': { 'train': 473, 'test': 473, }, 'mean': { 'R': 122.675, 'G': 116.669, 'B': 104.008, } }, }, 'cocostuff': { 'path_pytorch_model': 'data/models/deeplab_resnet101/cocostuff_rgb/deeplabv2_resnet101_VOC2012.pth', 'label_list': 'data/datasets/voc12/labels.txt', 'n_classes': 183, 'n_blocks': [3, 4, 23, 3], 'pyramids': [6, 3, 2, 1], 'image': { 'size': { 'train': 321, 'test': 513, }, 'mean': { 'R': 122.675, 'G': 116.669, 'B': 104.008, } }, }, }.get(dataset) cuda = cuda and torch.cuda.is_available() # Label list with open(CONFIG['label_list']) as f: classes = {} for label in f: label = label.rstrip().split('\t') classes[int(label[0])] = label[1].split(',')[0] # Load a model if model_path is None: state_dict = torch.load(CONFIG['path_pytorch_model']) else: state_dict = torch.load(model_path) # Model model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG['n_classes']) model.load_state_dict(state_dict) model.eval() if cuda: model.cuda() image_size = (CONFIG['image']['size']['test'], CONFIG['image']['size']['test']) # Image preprocessing image = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(float) image = cv2.resize(image, image_size) image_original = image.astype(np.uint8) image -= np.array([ float(CONFIG['image']['mean']['B']), float(CONFIG['image']['mean']['G']), float(CONFIG['image']['mean']['R']) ]) image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0) image = image.cuda() if cuda else image # Inference output = model(Variable(image, volatile=True)) output = F.upsample(output, size=image_size, mode='bilinear') output = F.softmax(output, dim=1) output = output.data.cpu().numpy()[0] if crf: output = dense_crf(image_original, output) labelmap = np.argmax(output.transpose(1, 2, 0), axis=2) labels = np.unique(labelmap) # Show results rows = np.floor(np.sqrt(len(labels) + 1)) cols = np.ceil((len(labels) + 1) / rows) plt.figure(figsize=(10, 10)) ax = plt.subplot(rows, cols, 1) ax.set_title('Input image') ax.imshow(image_original[:, :, ::-1]) ax.set_xticks([]) ax.set_yticks([]) for i, label in enumerate(labels): print '{0:3d}: {1}'.format(label, classes[label]) mask = labelmap == label ax = plt.subplot(rows, cols, i + 2) ax.set_title(classes[label]) ax.imshow(np.dstack((mask, ) * 3) * image_original[:, :, ::-1]) ax.set_xticks([]) ax.set_yticks([]) plt.show()
def main(dataset): WHITELIST = ["kernel_size", "stride", "padding", "dilation", "eps", "momentum"] CONFIG = { "voc12": { "path_caffe_model": "data/models/deeplab_resnet101/voc12/train2_iter_20000.caffemodel", "path_pytorch_model": "data/models/deeplab_resnet101/voc12/deeplabv2_resnet101_VOC2012.pth", "n_classes": 21, }, "init": { "path_caffe_model": "data/models/deeplab_resnet101/coco_init/init.caffemodel", "path_pytorch_model": "data/models/deeplab_resnet101/coco_init/deeplabv2_resnet101_COCO_init.pth", "n_classes": 91, }, }.get(dataset) params = parse_caffemodel(CONFIG["path_caffe_model"]) model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG["n_classes"]) model.eval() own_state = model.state_dict() rel_tol = 1e-7 state_dict = OrderedDict() for layer_name, layer_dict in params.items(): for param_name, values in layer_dict.items(): if param_name in WHITELIST and dataset != "coco_init" and dataset != "init": attribute = translate_layer_name(layer_name) attribute = eval("model." + attribute + "." + param_name) if isinstance(attribute, tuple): assert ( attribute[0] == values ), "Inconsistent values: {}@{}, {}@{}".format( attribute[0], translate_layer_name(layer_name) + "." + param_name, values, layer_name, ) else: assert ( abs(attribute - values) < rel_tol ), "Inconsistent values: {}@{}, {}@{}".format( attribute, translate_layer_name(layer_name) + "." + param_name, values, layer_name, ) print( layer_name.ljust(20), "->", param_name, attribute, values, ": Checked!", ) continue param_name = translate_layer_name(layer_name) + "." + param_name if param_name in own_state: values = torch.FloatTensor(values) values = values.view_as(own_state[param_name]) state_dict[param_name] = values print(layer_name.ljust(20), "->", param_name, ": Copied!") try: print("\033[32mVerify the converted model\033[00m") model.load_state_dict(state_dict) except: import traceback traceback.print_exc() print("\033[32mVerify with ignoring ASPP (strict=False)\033[00m") model.load_state_dict(state_dict, strict=False) print("Saving to", CONFIG["path_pytorch_model"]) torch.save(state_dict, CONFIG["path_pytorch_model"])
def main(config, cuda, gpu): # Configuration CONFIG = Dict(yaml.load(open(config))) # CUDA check cuda = cuda and torch.cuda.is_available() if cuda: gpu_ids = [int(string) for string in gpu.split(',')] current_device = torch.cuda.current_device() print('Running on', torch.cuda.get_device_name(current_device), gpu_ids) # Dataset dataset = CocoStuff10k( root=CONFIG.ROOT, split='train', image_size=513, crop_size=CONFIG.IMAGE.SIZE.TRAIN, scale=True, flip=True, ) # DataLoader loader = torch.utils.data.DataLoader( dataset=dataset, batch_size=CONFIG.BATCH_SIZE, num_workers=CONFIG.NUM_WORKERS, shuffle=True, ) loader_iter = iter(loader) # Model model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES) state_dict = torch.load(CONFIG.INIT_MODEL) model.load_state_dict(state_dict, strict=False) # Skip "aspp" layer model = nn.DataParallel(model, device_ids=gpu_ids) if cuda: model.cuda() # Optimizer optimizer = { 'sgd': torch.optim.SGD( # cf lr_mult and decay_mult in train.prototxt params=[{ 'params': get_lr_params(model.module, key='1x'), 'lr': CONFIG.LR, 'weight_decay': CONFIG.WEIGHT_DECAY }, { 'params': get_lr_params(model.module, key='10x'), 'lr': 10 * CONFIG.LR, 'weight_decay': CONFIG.WEIGHT_DECAY }, { 'params': get_lr_params(model.module, key='20x'), 'lr': 20 * CONFIG.LR, 'weight_decay': 0.0 }], momentum=CONFIG.MOMENTUM, ), }.get(CONFIG.OPTIMIZER) # Loss definition criterion = CrossEntropyLoss2d(ignore_index=CONFIG.IGNORE_LABEL) if cuda: criterion.cuda() # TensorBoard Logger writer = SummaryWriter(CONFIG.LOG_DIR) loss_meter = MovingAverageValueMeter(20) model.train() model.module.scale.freeze_bn() for iteration in tqdm( range(1, CONFIG.ITER_MAX + 1), total=CONFIG.ITER_MAX, leave=False, dynamic_ncols=True, ): # Set a learning rate poly_lr_scheduler( optimizer=optimizer, init_lr=CONFIG.LR, iter=iteration - 1, lr_decay_iter=CONFIG.LR_DECAY, max_iter=CONFIG.ITER_MAX, power=CONFIG.POLY_POWER, ) # Clear gradients (ready to accumulate) optimizer.zero_grad() iter_loss = 0 for i in range(1, CONFIG.ITER_SIZE + 1): data, target = next(loader_iter) # Image data = data.cuda() if cuda else data data = Variable(data) # Propagate forward outputs = model(data) # Loss loss = 0 for output in outputs: # Resize target for {100%, 75%, 50%, Max} outputs target_ = resize_target(target, output.size(2)) target_ = target_.cuda() if cuda else target_ target_ = Variable(target_) # Compute crossentropy loss loss += criterion(output, target_) # Backpropagate (just compute gradients wrt the loss) loss /= float(CONFIG.ITER_SIZE) loss.backward() iter_loss += loss.data[0] # Reload dataloader if ((iteration - 1) * CONFIG.ITER_SIZE + i) % len(loader) == 0: loader_iter = iter(loader) loss_meter.add(iter_loss) # Update weights with accumulated gradients optimizer.step() # TensorBoard if iteration % CONFIG.ITER_TF == 0: writer.add_scalar('train_loss', loss_meter.value()[0], iteration) for i, o in enumerate(optimizer.param_groups): writer.add_scalar('train_lr_group{}'.format(i), o['lr'], iteration) if iteration % 1000 != 0: continue for name, param in model.named_parameters(): name = name.replace('.', '/') writer.add_histogram(name, param, iteration, bins="auto") if param.requires_grad: writer.add_histogram(name + '/grad', param.grad, iteration, bins="auto") # Save a model if iteration % CONFIG.ITER_SNAP == 0: torch.save( model.module.state_dict(), osp.join(CONFIG.SAVE_DIR, 'checkpoint_{}.pth'.format(iteration)), ) # Save a model if iteration % 100 == 0: torch.save( model.module.state_dict(), osp.join(CONFIG.SAVE_DIR, 'checkpoint_current.pth'), ) torch.save( model.module.state_dict(), osp.join(CONFIG.SAVE_DIR, 'checkpoint_final.pth'), )
def main(config, embedding, model_path, run, imagedataset, local_rank, resnet, bkg): rank, world_size, device_id, device = setup(local_rank) print("Local rank: {} Rank: {} World Size: {} Device_id: {} Device: {}". format(local_rank, rank, world_size, device_id, device)) pth_extn = '.pth.tar' # Configuration CONFIG = Dict(yaml.load(open(config))) datadir = os.path.join('data/datasets', imagedataset) print("Split dir: ", datadir) savedir = osp.dirname(model_path) epoch = re.findall("checkpoint_(.*)\." + pth_extn[1:], osp.basename(model_path))[-1] if run == 'zlss' or run == 'flss': val = np.load(datadir + '/split/test_list.npy') visible_classes = np.load(datadir + '/split/novel_cls.npy') if bkg: visible_classes = np.asarray(np.concatenate( [np.array([0]), visible_classes]), dtype=int) elif run == 'gzlss' or run == 'gflss': val = np.load(datadir + '/split/test_list.npy') vals_cls = np.asarray(np.concatenate([ np.load(datadir + '/split/seen_cls.npy'), np.load(datadir + '/split/val_cls.npy') ]), dtype=int) if bkg: vals_cls = np.asarray(np.concatenate([np.array([0]), vals_cls]), dtype=int) valu_cls = np.load(datadir + '/split/novel_cls.npy') visible_classes = np.concatenate([vals_cls, valu_cls]) else: print("invalid run ", run) sys.exit() cls_map = np.array([255] * 256) for i, n in enumerate(visible_classes): cls_map[n] = i if run == 'gzlss' or run == 'gflss': novel_cls_map = np.array([255] * 256) for i, n in enumerate(list(valu_cls)): novel_cls_map[cls_map[n]] = i seen_cls_map = np.array([255] * 256) for i, n in enumerate(list(vals_cls)): seen_cls_map[cls_map[n]] = i all_labels = np.genfromtxt(datadir + '/labels_2.txt', delimiter='\t', usecols=1, dtype='str') print("Visible Classes: ", visible_classes) # Dataset dataset = get_dataset(CONFIG.DATASET)( train=None, test=val, root=CONFIG.ROOT, split=CONFIG.SPLIT.TEST, base_size=CONFIG.IMAGE.SIZE.TEST, mean=(CONFIG.IMAGE.MEAN.B, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.R), warp=CONFIG.WARP_IMAGE, scale=None, flip=False, ) random.seed(42) if embedding == 'word2vec': class_emb = pickle.load( open(datadir + '/word_vectors/word2vec.pkl', "rb")) elif embedding == 'fasttext': class_emb = pickle.load( open(datadir + '/word_vectors/fasttext.pkl', "rb")) elif embedding == 'fastnvec': class_emb = np.concatenate([ pickle.load(open(datadir + '/word_vectors/fasttext.pkl', "rb")), pickle.load(open(datadir + '/word_vectors/word2vec.pkl', "rb")) ], axis=1) else: print("invalid emb ", embedding) sys.exit() class_emb = class_emb[visible_classes] class_emb = F.normalize(torch.tensor(class_emb), p=2, dim=1).cuda() print("Embedding dim: ", class_emb.shape[1]) print("# Visible Classes: ", class_emb.shape[0]) # DataLoader loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=CONFIG.BATCH_SIZE.TEST, num_workers=CONFIG.NUM_WORKERS, shuffle=False, sampler=DistributedSampler( dataset, num_replicas=world_size, rank=rank, shuffle=False), pin_memory=True, drop_last=True) torch.set_grad_enabled(False) # Model model = DeepLabV2_ResNet101_MSC(class_emb.shape[1], class_emb, resnet=resnet) state_dict = torch.load(model_path, map_location='cpu') model = DistributedDataParallel(model.to(device), device_ids=[rank]) new_state_dict = OrderedDict() if resnet == 'spnet': for k, v in state_dict['state_dict'].items(): name = k.replace("scale", "base") # 'scale'->base name = name.replace("stages.", "") new_state_dict[name] = v else: new_state_dict = state_dict['state_dict'] model.load_state_dict(new_state_dict) del state_dict model.eval() targets, outputs = [], [] loader_iter = iter(loader) iterations = len(loader_iter) print("Iterations: {}".format(iterations)) pbar = tqdm(loader, total=iterations, leave=False, dynamic_ncols=True, position=rank) for iteration in pbar: data, target, img_id = next(loader_iter) # Image data = data.to(device) # Forward propagation output = model(data) output = F.interpolate(output, size=data.shape[2:], mode="bilinear", align_corners=False) output = F.softmax(output, dim=1) target = cls_map[target.numpy()] remote_target = torch.tensor(target).to(device) if rank == 0: remote_target = torch.zeros_like(remote_target).to(device) output = torch.argmax(output, dim=1).cpu().numpy() remote_output = torch.tensor(output).to(device) if rank == 0: remote_output = torch.zeros_like(remote_output).to(device) for o, t in zip(output, target): outputs.append(o) targets.append(t) torch.distributed.reduce(remote_output, dst=0) torch.distributed.reduce(remote_target, dst=0) torch.distributed.barrier() if rank == 0: remote_output = remote_output.cpu().numpy() remote_target = remote_target.cpu().numpy() for o, t in zip(remote_output, remote_target): outputs.append(o) targets.append(t) if rank == 0: if run == 'gzlss' or run == 'gflss': score, class_iou = scores_gzsl(targets, outputs, n_class=len(visible_classes), seen_cls=cls_map[vals_cls], unseen_cls=cls_map[valu_cls]) else: score, class_iou = scores(targets, outputs, n_class=len(visible_classes)) for k, v in score.items(): print(k, v) score["Class IoU"] = {} for i in range(len(visible_classes)): score["Class IoU"][all_labels[visible_classes[i]]] = class_iou[i] name = "" name = model_path.replace(pth_extn, "_" + run + ".json") if bkg == True: with open(name.replace('.json', '_bkg.json'), "w") as f: json.dump(score, f, indent=4, sort_keys=True) else: with open(name, "w") as f: json.dump(score, f, indent=4, sort_keys=True) print(score["Class IoU"]) return