import loadseg import expdir from indexdata import load_image_to_label from labelprobe import cached_memmap from labelprobe_pytorch import get_seg_size, iou_intersect_d, iou_union_d from linearprobe_pytorch import CustomLayer import numpy as np import torch import torch.nn as nn from torch.autograd import Variable import os import time def probe_linear(directory, blob, suffix='', start=None, end=None,batch_size=16, ahead=4, quantile=0.005, bias=False, positive=False, cuda=False, force=False): qcode = ('%f' % quantile).replace('0.','.').rstrip('0') ed = expdir.ExperimentDirectory(directory) if (ed.has_mmap(blob=blob, part='linear_ind_ious%s' % suffix) and ed.has_mmap(blob=blob, part='linear_set_ious%s' % suffix)): print('Linear weights have already been probed.') print ed.mmap_filename(blob=blob, part='linear_set_val_ious%s' % suffix) if not force: return else: print('Forcefully continuing...') info = ed.load_info() seg_size = get_seg_size(info.input_dim) blob_info = ed.load_info(blob=blob) ds = loadseg.SegmentationData(info.dataset) shape = blob_info.shape N = shape[0] # number of total images K = shape[1] # number of units in given blob L = ds.label_size() # number of labels if quantile == 1: thresh = np.zeros((K,1,1)) else: quantdata = ed.open_mmap(blob=blob, part='quant-*', shape=(K, -1)) threshold = quantdata[:, int(round(quantdata.shape[1] * quantile))] thresh = threshold[:, np.newaxis, np.newaxis] fn_read = ed.mmap_filename(blob=blob) blobdata = cached_memmap(fn_read, mode='r', dtype='float32', shape=shape) image_to_label = load_image_to_label(directory) if ed.has_mmap(blob=blob, part='linear_ind_ious%s' % suffix, inc=True) and not force: assert(ed.has_mmap(blob=blob, part='linear_set_ious%s' % suffix, inc=True)) ind_ious = ed.open_mmap(blob=blob, part='linear_ind_ious%s' % suffix, mode='r+', inc=True, dtype='float32', shape=(L,N)) set_ious = ed.open_mmap(blob=blob, part='linear_set_ious%s' % suffix, mode='r+', inc=True, dtype='float32', shape=(L,)) set_ious_train = ed.open_mmap(blob=blob, part='linear_set_train_ious%s' % suffix, mode='r+', inc=True, dtype='float32', shape=(L,)) try: set_ious_val = ed.open_mmap(blob=blob, part='linear_set_val_ious%s' % suffix, mode='r+', inc=True, dtype='float32', shape=(L,)) except: set_ious_val = ed.open_mmap(blob=blob, part='linear_set_val_ious%s' % suffix, mode='r+', dtype='float32', shape=(L,)) else: ind_ious = ed.open_mmap(blob=blob, part='linear_ind_ious%s' % suffix, mode='w+', dtype='float32', shape=(L,N)) set_ious = ed.open_mmap(blob=blob, part='linear_set_ious%s' % suffix, mode='w+', dtype='float32', shape=(L,)) set_ious_train = ed.open_mmap(blob=blob, part='linear_set_train_ious%s' % suffix, mode='w+', dtype='float32', shape=(L,)) set_ious_val = ed.open_mmap(blob=blob, part='linear_set_val_ious%s' % suffix, mode='w+', dtype='float32', shape=(L,)) if start is None: start = 1 if end is None: end = L for label_i in range(start, end): if ed.has_mmap(blob=blob, part='label_i_%d_weights%s' % (label_i, suffix)): try: weights = ed.open_mmap(blob=blob, part='label_i_%d_weights%s' % (label_i, suffix),mode='r', dtype='float32', shape=(K,)) except ValueError: # SUPPORTING LEGACY CODE (TODO: Remove) weights = ed.open_mmap(blob=blob, part='label_i_%d_weights%s' % (label_i, suffix), mode='r', dtype=float, shape=(K,)) elif ed.has_mmap(blob=blob, part='linear_weights%s' % suffix): all_weights = ed.open_mmap(blob=blob, part='linear_weights%s' % suffix, mode='r', dtype='float32', shape=(L,K)) weights = all_weights[label_i] if not np.any(weights): print('Label %d does not have associated weights to it, so skipping.' % label_i) continue else: print('Label %d does not have associated weights to it, so skipping.' % label_i) continue if bias: if ed.has_mmap(blob=blob, part='label_i_%d_bias%s' % (label_i, suffix)): bias_v = ed.open_mmap(blob=blob, part='label_i_%d_bias%s' % (label_i, suffix), mode='r', dtype='float32', shape=(1,)) else: assert(ed.has_mmap(blob=blob, part='linear_bias%s' % suffix)) all_bias_v = ed.open_mmap(blob=blob, part='linear_bias%s' % suffix, mode='r', dtype='float32', shape=(L,)) bias_v = np.array([all_bias_v[label_i]]) label_categories = ds.label[label_i]['category'].keys() label_name = ds.name(category=None, j=label_i) label_idx = np.where(image_to_label[:, label_i])[0] loader = loadseg.SegmentationPrefetcher(ds, categories=label_categories, indexes=label_idx, once=True, batch_size=batch_size, ahead=ahead, thread=True) loader_idx = loader.indexes num_imgs = len(loader.indexes) print('Probing with learned weights for label %d (%s) with %d images...' % ( label_i, label_name, num_imgs)) model = CustomLayer(K, upsample=True, up_size=seg_size, act=True, bias=bias, positive=positive, cuda=cuda) model.weight.data[...] = torch.Tensor(weights) if bias: model.bias.data[...] = torch.Tensor(bias_v) if cuda: model.cuda() model.eval() iou_intersects = np.zeros(num_imgs) iou_unions = np.zeros(num_imgs) i = 0 for batch in loader.batches(): start_t = time.time() if (i+1)*batch_size < num_imgs: idx = range(i*batch_size, (i+1)*batch_size) else: idx = range(i*batch_size, num_imgs) i += 1 input = torch.Tensor((blobdata[loader_idx[idx]] > thresh).astype(float)) input_var = (Variable(input.cuda(), volatile=True) if cuda else Variable(input, volatile=True)) target_ = [] for rec in batch: for cat in label_categories: if rec[cat] != []: if type(rec[cat]) is np.ndarray: target_.append(np.max((rec[cat] == label_i).astype(float), axis=0)) else: target_.append(np.ones(seg_size)) break target = torch.Tensor(target_) target_var = (Variable(target.cuda(), volatile=True) if cuda else Variable(target, volatile=True)) #target_var = Variable(target.unsqueeze(1).expand_as( # input_var).cuda() if cuda # else target.unsqueeze(1).expand_as(input_var)) output_var = model(input_var) iou_intersects[idx] = np.squeeze(iou_intersect_d(output_var, target_var).data.cpu().numpy()) iou_unions[idx] = np.squeeze(iou_union_d(output_var, target_var).data.cpu().numpy()) print('Batch: %d/%d\tTime: %f secs\tAvg Ind IOU: %f' % (i, num_imgs/batch_size, time.time()-start_t, np.mean(np.true_divide(iou_intersects[idx], iou_unions[idx] + 1e-20)))) loader.close() label_ind_ious = np.true_divide(iou_intersects, iou_unions + 1e-20) label_set_iou = np.true_divide(np.sum(iou_intersects), np.sum(iou_unions) + 1e-20) ind_ious[label_i, loader_idx] = label_ind_ious set_ious[label_i] = label_set_iou train_idx = [i for i in range(len(loader_idx)) if ds.split(loader_idx[i]) == 'train'] val_idx = [i for i in range(len(loader_idx)) if ds.split(loader_idx[i]) == 'val'] set_ious_train[label_i] = np.true_divide(np.sum(iou_intersects[train_idx]), np.sum(iou_unions[train_idx]) + 1e-20) set_ious_val[label_i] = np.true_divide(np.sum(iou_intersects[val_idx]), np.sum(iou_unions[val_idx]) + 1e-20) print('Label %d (%s) Set IOU: %f, Train Set IOU: %f, Val Set IOU: %f, Max Ind IOU: %f' % (label_i, label_name, label_set_iou, set_ious_train[label_i], set_ious_val[label_i], np.max(label_ind_ious))) ed.finish_mmap(ind_ious) ed.finish_mmap(set_ious) ed.finish_mmap(set_ious_train) ed.finish_mmap(set_ious_val) if __name__ == '__main__': import argparse import sys import traceback try: parser = argparse.ArgumentParser() parser.add_argument('--directory', default='.') parser.add_argument('--blobs', nargs='*') parser.add_argument('--suffix', type=str, default='') parser.add_argument('--start', type=int, default=1) parser.add_argument('--end', type=int, default=1198) parser.add_argument('--batch_size', type=int, default=16) parser.add_argument('--ahead', type=int, default=4) parser.add_argument('--quantile', type=float, default=0.005) parser.add_argument('--bias', action='store_true', default=False) parser.add_argument('--positive', action='store_true', default=False) parser.add_argument('--force', action='store_true', default=False) parser.add_argument('--num_filters', type=int, nargs='*', default=None) parser.add_argument('--gpu', type=int, default=None) args = parser.parse_args() gpu = args.gpu cuda = True if gpu is not None else False use_mult_gpu = isinstance(gpu, list) if cuda: if use_mult_gpu: os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu).strip('[').strip(']') else: os.environ['CUDA_VISIBLE_DEVICES'] = '%d' % gpu print(torch.cuda.device_count(), use_mult_gpu, cuda) for blob in args.blobs: if args.num_filters is not None: suffixes = ['%s_num_filters_%d' % (args.suffix, n) for n in args.num_filters] else: suffixes = [args.suffix] for i in range(len(suffixes)): probe_linear(args.directory, blob, suffix=suffixes[i], start=args.start, end=args.end, batch_size=args.batch_size, ahead=args.ahead, quantile=args.quantile, bias=args.bias, positive=args.positive, cuda=cuda, force=args.force) except: traceback.print_exc(file=sys.stdout) sys.exit(1)