def get_iou_callable(): """ Get a pairwise box iou callable. """ with tf.Graph().as_default(), tf.device('/cpu:0'): A = tf.placeholder(tf.float32, shape=[None, 4]) B = tf.placeholder(tf.float32, shape=[None, 4]) iou = pairwise_iou(A, B) sess = tf.Session(config=get_default_sess_config()) return sess.make_callable(iou, [A, B])
def get_tf_nms(): """ Get a NMS callable. """ boxes = tf.placeholder(tf.float32, shape=[None, 4]) scores = tf.placeholder(tf.float32, shape=[None]) indices = tf.image.non_max_suppression(boxes, scores, config.RESULTS_PER_IM, config.FASTRCNN_NMS_THRESH) sess = tf.Session(config=get_default_sess_config()) return sess.make_callable(indices, [boxes, scores])
def get_iou_callable(): """ Get a pairwise box iou callable. """ # We don't want the dataflow process to touch CUDA # Data needs tensorflow. As a result, the training cannot run on GPUs with # EXCLUSIVE_PROCESS mode, unless you disable multiprocessing prefetch. with tf.Graph().as_default(), tf.device('/cpu:0'): A = tf.placeholder(tf.float32, shape=[None, 4]) B = tf.placeholder(tf.float32, shape=[None, 4]) iou = pairwise_iou(A, B) sess = tf.Session(config=get_default_sess_config()) return sess.make_callable(iou, [A, B])
if __name__ == '__main__': import argparse from tensorpack.dataflow import TestDataSpeed from tensorpack.tfutils import get_default_sess_config parser = argparse.ArgumentParser() parser.add_argument('--data', required=True) parser.add_argument('--batch', type=int, default=32) parser.add_argument('--aug', choices=['train', 'val'], default='val') parser.add_argument('--symbolic', action='store_true') args = parser.parse_args() if not args.symbolic: augs = fbresnet_augmentor(args.aug == 'train') df = get_imagenet_dataflow(args.data, 'train', args.batch, augs) # For val augmentor, Should get >100 it/s (i.e. 3k im/s) here on a decent E5 server. TestDataSpeed(df).start() else: assert args.aug == 'train' data = get_imagenet_tfdata(args.data, 'train', args.batch) itr = data.make_initializable_iterator() dp = itr.get_next() dpop = tf.group(*dp) with tf.Session(config=get_default_sess_config()) as sess: sess.run(itr.initializer) for _ in tqdm.trange(200): sess.run(dpop) for _ in tqdm.trange(5000, smoothing=0.1): sess.run(dpop)
shape=[num_train_images, 128], trainable=False) net = ResNetModel(num_output=(2048, 128) if args.v2 else (128, )) with TowerContext("", is_training=False): feat = net.forward(image_input) feat = tf.math.l2_normalize(feat, axis=1) # Nx128 all_feat = hvd.allgather(feat) # GN x 128 all_idx_input = hvd.allgather(idx_input) # GN update_buffer = tf.scatter_update(feat_buffer, all_idx_input, all_feat) dist = tf.matmul(feat, tf.transpose(feat_buffer)) # N x #DS _, topk_indices = tf.math.top_k(dist, k=args.top_k) # Nxtopk train_ds = build_dataflow(local_train_files) config = get_default_sess_config() config.gpu_options.visible_device_list = str(hvd.local_rank()) def evaluate(checkpoint_file): result_file = get_checkpoint_path( checkpoint_file) + f".knn{args.top_k}.txt" if os.path.isfile(result_file): logger.info(f"Skipping evaluation of {result_file}.") return with tf.Session(config=config) as sess: sess.run(tf.global_variables_initializer()) SmartInit(checkpoint_file).init(sess) for batch_img, batch_idx in tqdm.tqdm(train_ds, total=len(train_ds)): sess.run(update_buffer, feed_dict={