示例#1
0
def get_data(name, data_dir, meta_dir, gpu_nums):
    isTrain = name == 'train'
    ds = PascalVOC12(data_dir, meta_dir, name, shuffle=True)


    if isTrain:#special augmentation
        shape_aug = [RandomResize(xrange=(0.7, 1.5), yrange=(0.7, 1.5),
                            aspect_ratio_thres=0.15),
                     RandomCropWithPadding(args.crop_size,IGNORE_LABEL),
                     Flip(horiz=True),
                     ]
    else:
        shape_aug = []

    ds = AugmentImageComponents(ds, shape_aug, (0, 1), copy=False)

    def f(ds):
        image, label = ds
        m = np.array([104, 116, 122])
        const_arr = np.resize(m, (1,1,3))  # NCHW
        image = image - const_arr
        return image, label

    ds = MapData(ds, f)
    if isTrain:
        ds = BatchData(ds, args.batch_size*gpu_nums)
        ds = PrefetchDataZMQ(ds, 1)
    else:
        ds = BatchData(ds, 1)
    return ds
示例#2
0
def get_mnist_data(is_train, image_size, batchsize):
    ds = MNISTCh('train' if is_train else 'test', shuffle=True)

    if is_train:
        augs = [
            imgaug.RandomApplyAug(imgaug.RandomResize((0.8, 1.2), (0.8, 1.2)),
                                  0.3),
            imgaug.RandomApplyAug(imgaug.RotationAndCropValid(15), 0.5),
            imgaug.RandomApplyAug(
                imgaug.SaltPepperNoise(white_prob=0.01, black_prob=0.01),
                0.25),
            imgaug.Resize((224, 224), cv2.INTER_AREA)
        ]
        ds = AugmentImageComponent(ds, augs)
        ds = PrefetchData(ds, 128 * 10, multiprocessing.cpu_count())
        ds = BatchData(ds, batchsize)
        ds = PrefetchData(ds, 256, 4)
    else:
        # no augmentation, only resizing
        augs = [
            imgaug.Resize((image_size, image_size), cv2.INTER_CUBIC),
        ]
        ds = AugmentImageComponent(ds, augs)
        ds = BatchData(ds, batchsize)
        ds = PrefetchData(ds, 20, 2)
    return ds
示例#3
0
def get_data(name, data_dir, meta_dir, gpu_nums):
    isTrain = True if 'train' in name else False
    ds = Camvid(data_dir, meta_dir, name, shuffle=True)


    if isTrain:
        ds = MapData(ds, RandomResize)

    if isTrain:
        shape_aug = [
                     RandomCropWithPadding(args.crop_size,IGNORE_LABEL),
                     Flip(horiz=True),
                     ]
    else:
        shape_aug = []

    ds = AugmentImageComponents(ds, shape_aug, (0, 1), copy=False)

    def f(ds):
        image, label = ds
        m = np.array([104, 116, 122])
        const_arr = np.resize(m, (1,1,3))  # NCHW
        image = image - const_arr
        return image, label

    ds = MapData(ds, f)
    if isTrain:
        ds = BatchData(ds, args.batch_size*gpu_nums)
        ds = PrefetchDataZMQ(ds, 1)
    else:
        ds = BatchData(ds, 1)
    return ds
示例#4
0
    def get_input_flow(self):
        ds_train = CellImageDataManagerTrain()
        # ds_train = MapDataComponent(ds_train, random_affine)  # TODO : no improvement?
        ds_train = MapDataComponent(ds_train, random_color)
        # ds_train = MapDataComponent(ds_train, random_scaling)
        ds_train = MapDataComponent(
            ds_train,
            mask_size_normalize)  # Resize by instance size - normalization
        ds_train = MapDataComponent(
            ds_train, lambda x: resize_shortedge_if_small(x, self.img_size))
        ds_train = MapDataComponent(
            ds_train, lambda x: random_crop(x, self.img_size, self.img_size))
        ds_train = MapDataComponent(ds_train, random_flip_lr)
        ds_train = MapDataComponent(ds_train, random_flip_ud)
        # ds_train = MapDataComponent(ds_train, data_to_elastic_transform_wrapper)
        ds_train = MapDataComponent(ds_train, erosion_mask)
        ds_train = MapData(
            ds_train, lambda x: data_to_segment_input(
                x, is_gray=False, unet_weight=True))
        ds_train = PrefetchData(ds_train, 256, 24)
        ds_train = BatchData(ds_train, self.batchsize)
        ds_train = MapDataComponent(ds_train, data_to_normalize1)

        ds_valid = CellImageDataManagerValid()
        ds_valid = MapDataComponent(
            ds_valid, lambda x: resize_shortedge_if_small(x, self.img_size))
        ds_valid = MapDataComponent(
            ds_valid, lambda x: random_crop(x, self.img_size, self.img_size))
        ds_valid = MapDataComponent(ds_valid, erosion_mask)
        ds_valid = MapData(
            ds_valid, lambda x: data_to_segment_input(
                x, is_gray=False, unet_weight=True))
        ds_valid = PrefetchData(ds_valid, 20, 12)
        ds_valid = BatchData(ds_valid, self.batchsize, remainder=True)
        ds_valid = MapDataComponent(ds_valid, data_to_normalize1)

        ds_valid2 = CellImageDataManagerValid()
        ds_valid2 = MapDataComponent(
            ds_valid2, lambda x: resize_shortedge_if_small(x, self.img_size))
        ds_valid2 = MapDataComponent(
            ds_valid2,
            lambda x: center_crop_if_tcga(x, self.img_size, self.img_size))
        # ds_valid2 = MapDataComponent(ds_valid2, lambda x: resize_shortedge(x, self.img_size))
        ds_valid2 = MapData(ds_valid2,
                            lambda x: data_to_segment_input(x, is_gray=False))
        ds_valid2 = MapDataComponent(ds_valid2, data_to_normalize1)

        ds_test = CellImageDataManagerTest()
        ds_test = MapDataComponent(
            ds_test, lambda x: resize_shortedge_if_small(x, self.img_size))
        # ds_test = MapDataComponent(ds_test, lambda x: resize_shortedge(x, self.img_size))
        ds_test = MapData(ds_test, lambda x: data_to_image(x, is_gray=False))
        ds_test = MapDataComponent(ds_test, data_to_normalize1)

        return ds_train, ds_valid, ds_valid2, ds_test
示例#5
0
def batch_dataflow(df, batch_size):
    """
    The function builds batch dataflow from the input dataflow of samples

    :param df: dataflow of samples
    :param batch_size: batch size
    :return: dataflow of batches
    """
    df = BatchData(df, batch_size, use_list=False)
    df = MapData(df, lambda x: ([x[0]], [x[2]]))
    df.reset_state()
    return df
示例#6
0
def get_data(name, meta_dir, gpu_nums):
    isTrain = True if 'train' in name else False

    m = np.array([104, 116, 122])
    const_arr = np.resize(m, (1, 1, 3))  # NCHW
    const_arr = np.zeros(
        (args.crop_size[0], args.crop_size[1], 3)) + const_arr  #broadcast

    if isTrain:
        #ds = FakeData([[1024, 2048, 3], [ 1024, 2048]], 5000, random=False, dtype='uint8')
        #ds = FakeData([[CROP_HEIGHT, CROP_HEIGHT, 3], [CROP_HEIGHT, CROP_HEIGHT]], 5000,random=False, dtype='uint8')
        ds = CityscapesFiles(base_dir, meta_dir, name, shuffle=True)
        parallel = min(3, multiprocessing.cpu_count())
        augmentors = [
            RandomCropWithPadding(args.crop_size),
            Flip(horiz=True),
        ]
        aug = imgaug.AugmentorList(augmentors)

        def mapf(ds):
            img, label = ds
            img = cv2.imread(img, cv2.IMREAD_COLOR)
            label = cv2.imread(label, cv2.IMREAD_GRAYSCALE)
            img, params = aug.augment_return_params(img)
            label = aug._augment(label, params)
            img = img - const_arr  # very time-consuming
            return img, label

        #ds = MapData(ds, mapf)
        ds = MultiThreadMapData(ds,
                                parallel,
                                mapf,
                                buffer_size=500,
                                strict=True)
        #ds = MapData(ds, reduce_mean_rgb)

        ds = BatchData(ds, args.batch_size * gpu_nums)
        #ds = PrefetchDataZMQ(ds, 1)
    else:

        def imgread(ds):
            img, label = ds
            img = cv2.imread(img, cv2.IMREAD_COLOR)
            label = cv2.imread(label, cv2.IMREAD_GRAYSCALE)
            return [img, label]

        ds = CityscapesFiles(base_dir, meta_dir, name, shuffle=False)
        ds = MapData(ds, imgread)
        ds = BatchData(ds, 1)

    return ds
    def get_input_flow(self):
        ds_train = CellImageDataManagerTrain()
        # Augmentation :
        ds_train = MapDataComponent(ds_train, random_affine)
        ds_train = MapDataComponent(ds_train, random_color)
        # ds_train = MapDataComponent(ds_train, random_color2)  # not good
        ds_train = MapDataComponent(ds_train, random_scaling)
        ds_train = MapDataComponent(
            ds_train, lambda x: resize_shortedge_if_small(x, 224))
        ds_train = MapDataComponent(ds_train,
                                    lambda x: random_crop(x, 224, 224))
        ds_train = MapDataComponent(ds_train, random_flip_lr)
        # ds_train = MapDataComponent(ds_train, data_to_elastic_transform_wrapper)
        ds_train = MapDataComponent(ds_train, random_flip_ud)
        if self.unet_weight:
            ds_train = MapDataComponent(ds_train, erosion_mask)
        ds_train = PrefetchData(ds_train, 1000, 24)
        ds_train = MapData(
            ds_train, lambda x: data_to_segment_input(x, not self.is_color,
                                                      self.unet_weight))
        ds_train = BatchData(ds_train, self.batchsize)
        ds_train = MapDataComponent(ds_train, data_to_normalize1)
        ds_train = PrefetchData(ds_train, 10, 2)

        ds_valid = CellImageDataManagerValid()
        ds_valid = MapDataComponent(ds_valid,
                                    lambda x: center_crop(x, 224, 224))
        if self.unet_weight:
            ds_valid = MapDataComponent(ds_valid, erosion_mask)
        ds_valid = MapData(
            ds_valid, lambda x: data_to_segment_input(x, not self.is_color,
                                                      self.unet_weight))
        ds_valid = BatchData(ds_valid, self.batchsize, remainder=True)
        ds_valid = MapDataComponent(ds_valid, data_to_normalize1)
        ds_valid = PrefetchData(ds_valid, 20, 24)

        ds_valid2 = CellImageDataManagerValid()
        ds_valid2 = MapDataComponent(
            ds_valid2, lambda x: resize_shortedge_if_small(x, 224))
        ds_valid2 = MapData(
            ds_valid2, lambda x: data_to_segment_input(x, not self.is_color))
        ds_valid2 = MapDataComponent(ds_valid2, data_to_normalize1)

        ds_test = CellImageDataManagerTest()
        ds_test = MapDataComponent(ds_test,
                                   lambda x: resize_shortedge_if_small(x, 224))
        ds_test = MapData(ds_test,
                          lambda x: data_to_image(x, not self.is_color))
        ds_test = MapDataComponent(ds_test, data_to_normalize1)

        return ds_train, ds_valid, ds_valid2, ds_test
示例#8
0
def get_dataflow_batch(is_train, batchsize):
    ds = get_dataflow(is_train)
    ds = PrefetchData(ds, 1000, multiprocessing.cpu_count())
    ds = BatchData(ds, batchsize)
    ds = PrefetchData(ds, 10, 4)

    return ds
示例#9
0
def get_remote_dataflow(port, nr_prefetch=1000, nr_thread=1):
    ipc = 'ipc:///tmp/ipc-socket'
    tcp = 'tcp://0.0.0.0:%d' % port
    data_loader = RemoteDataZMQ(ipc, tcp, hwm=10000)
    data_loader = BatchData(data_loader, batch_size=hp.train.batch_size)
    data_loader = PrefetchData(data_loader, nr_prefetch, nr_thread)
    return data_loader
示例#10
0
文件: train.py 项目: qq456cvb/UKPGAN
def main(cfg):
    print(cfg)
    
    tf.reset_default_graph()
    
    logger.set_logger_dir('tflogs', action='d')

    copyfile(hydra.utils.to_absolute_path('model.py'), 'model.py')
    copyfile(hydra.utils.to_absolute_path('dataflow.py'), 'dataflow.py')
    
    if cfg.cat_name == 'smpl':
        train_df = SMPLDataFlow(cfg, True, 1000)
        val_df = VisSMPLDataFlow(cfg, True, 1000, port=1080)
    else:
        train_df = ShapeNetDataFlow(cfg, cfg.data.train_txt, True)
        val_df = VisDataFlow(cfg, cfg.data.val_txt, False, port=1080)
    
    config = TrainConfig(
        model=Model(cfg),
        dataflow=BatchData(PrefetchData(train_df, cpu_count() // 2, cpu_count() // 2), cfg.batch_size),
        callbacks=[
            ModelSaver(),
            SimpleMovingAverage(['recon_loss', 'GAN/loss_d', 'GAN/loss_g', 'GAN/gp_loss', 'symmetry_loss'], 100),
            PeriodicTrigger(val_df, every_k_steps=30)
        ],
        monitors=tensorpack.train.DEFAULT_MONITORS() + [ScalarPrinter(enable_step=True, enable_epoch=False)],
        max_epoch=10
    )
    launch_train_with_config(config, SimpleTrainer())
def get_dataflow(annot_path, img_dir, batch_size):
    """
    This function initializes the tensorpack dataflow and serves generator
    for training operation.

    :param annot_path: path to the annotation file
    :param img_dir: path to the images
    :param batch_size: batch size
    :return: dataflow object
    """
    df = CocoDataFlow((368, 368), annot_path, img_dir)
    df.prepare()
    df = MapData(df, read_img)
    df = MapData(df, gen_mask)
    df = MapData(df, augment)
    df = MapData(df, apply_mask)
    df = MapData(df, build_sample)
    df = PrefetchDataZMQ(df, nr_proc=4)  #df = PrefetchData(df, 2, 1)
    df = BatchData(df, batch_size, use_list=False)
    df = MapData(
        df, lambda x: ([x[0], x[1], x[2]], [
            x[3], x[4], x[3], x[4], x[3], x[4], x[3], x[4], x[3], x[4], x[3],
            x[4]
        ]))
    df.reset_state()

    return df
def get_default_dataflow_batch(batchsize=32):
    ds = get_default_dataflow()
    ds = MapData(ds, data_to_segment_input)
    ds = BatchData(ds, batchsize)
    ds = MapDataComponent(ds, data_to_normalize01)
    ds = PrefetchData(ds, 10, 2)

    return ds
示例#13
0
def get_dataflow_batch(path, is_train, batchsize, img_path=None):
    logger.info('dataflow img_path=%s' % img_path)
    ds = get_dataflow(path, is_train, img_path=img_path)
    ds = BatchData(ds, batchsize)
    if is_train:
        ds = PrefetchData(ds, batchsize*2, 1)
    else:
        ds = PrefetchData(ds, batchsize*2, 1)

    return ds
示例#14
0
def get_dataflow_batch(path, clothe_class, is_train, batchsize, img_path=None):
    logger.info('dataflow img_path=%s' % img_path)
    ds = get_dataflow(path, is_train, clothe_class, img_path=img_path)
    ds = BatchData(ds, batchsize)
    if is_train:
        ds = PrefetchData(ds, 10, 2)
    else:
        ds = PrefetchData(ds, 50, 2)

    return ds
示例#15
0
def get_dataflow_batch(path, is_train, batchsize, img_path=None):
    logger.info('dataflow img_path=%s' % img_path)
    ds = get_dataflow(path, is_train, img_path=img_path)
    ds = BatchData(ds, batchsize)
    # if is_train:
    #     ds = PrefetchData(ds, 10, 2)
    # else:
    #     ds = PrefetchData(ds, 50, 2)

    return ds
示例#16
0
def get_data(name, data_dir, meta_dir, gpu_nums):
    isTrain = True if 'train' in name else False

    def imgread(ds):
        img, label = ds
        img = cv2.imread(img, cv2.IMREAD_COLOR)
        label = cv2.imread(label, cv2.IMREAD_GRAYSCALE)
        return img, label

    if isTrain:
        #ds = LMDBData('/data2/dataset/cityscapes/cityscapes_train.lmdb', shuffle=True)
        #ds = FakeData([[batch_size, CROP_HEIGHT, CROP_HEIGHT, 3], [batch_size, CROP_HEIGHT, CROP_HEIGHT, 1]], 5000, random=False, dtype='uint8')
        ds = PascalVOC12Files(data_dir, meta_dir, name, shuffle=True)
        ds = MultiThreadMapData(ds,4,imgread, buffer_size= 2)
        #ds = PrefetchDataZMQ(MapData(ds, ImageDecode), 1) #imagedecode is heavy
        ds = MapData(ds, RandomResize)
    else:
        ds = PascalVOC12Files(data_dir, meta_dir, name, shuffle=False)
        ds = MultiThreadMapData(ds, 4, imgread, buffer_size= 2)

    if isTrain:
        shape_aug = [
                     RandomCropWithPadding(args.crop_size,IGNORE_LABEL),
                     Flip(horiz=True),
                     ]
        ds = AugmentImageComponents(ds, shape_aug, (0, 1), copy=False)

    def reduce_mean_rgb(ds):
        image, label = ds
        m = np.array([104, 116, 122])
        const_arr = np.resize(m, (1,1,3))  # NCHW
        image = image - const_arr
        return image, label

    def MxnetPrepare(ds):
        data, label = ds
        data = np.transpose(data, (0, 3, 1, 2))  # NCHW
        label = label[:, :, :, None]
        label = np.transpose(label, (0, 3, 1, 2))  # NCHW
        dl = [[mx.nd.array(data[args.batch_size * i:args.batch_size * (i + 1)])] for i in
              range(gpu_nums)]  # multi-gpu distribute data, time-consuming!!!
        ll = [[mx.nd.array(label[args.batch_size * i:args.batch_size * (i + 1)])] for i in
              range(gpu_nums)]
        return dl, ll

    #ds = MapData(ds, reduce_mean_rgb)
    ds = MultiThreadMapData(ds, 4, reduce_mean_rgb, buffer_size=2)

    if isTrain:
        ds = FastBatchData(ds, args.batch_size*gpu_nums)
        ds = MapData(ds, MxnetPrepare)
        #ds = PrefetchDataZMQ(ds, 1)
    else:
        ds = BatchData(ds, 1)
    return ds
def get_hand_dataflow_batch(is_train, batchsize, img_path=None):
    logger.info('dataflow img_path=%s' % img_path)
    ds = get_hand_dataflow(is_train, img_path=img_path)
    ds = BatchData(ds, batchsize)
    if is_train:
        # ds = PrefetchDataZMQ(ds, 10, 2)
        PrefetchDataZMQ(ds, 8)
    else:
        ds = PrefetchData(ds, 50, 2)

    return ds
示例#18
0
def get_ilsvrc_data_alexnet(is_train, image_size, batchsize, directory):
    if is_train:
        if not directory.startswith('/'):
            ds = ILSVRCTTenthTrain(directory)
        else:
            ds = ILSVRC12(directory, 'train')
        augs = [
            imgaug.RandomApplyAug(imgaug.RandomResize((0.9, 1.2), (0.9, 1.2)),
                                  0.7),
            imgaug.RandomApplyAug(imgaug.RotationAndCropValid(15), 0.7),
            imgaug.RandomApplyAug(
                imgaug.RandomChooseAug([
                    imgaug.SaltPepperNoise(white_prob=0.01, black_prob=0.01),
                    imgaug.RandomOrderAug([
                        imgaug.BrightnessScale((0.8, 1.2), clip=False),
                        imgaug.Contrast((0.8, 1.2), clip=False),
                        # imgaug.Saturation(0.4, rgb=True),
                    ]),
                ]),
                0.7),
            imgaug.Flip(horiz=True),
            imgaug.ResizeShortestEdge(256, cv2.INTER_CUBIC),
            imgaug.RandomCrop((224, 224)),
        ]
        ds = AugmentImageComponent(ds, augs)
        ds = PrefetchData(ds, 1000, multiprocessing.cpu_count())
        ds = BatchData(ds, batchsize)
        ds = PrefetchData(ds, 10, 4)
    else:
        if not directory.startswith('/'):
            ds = ILSVRCTenthValid(directory)
        else:
            ds = ILSVRC12(directory, 'val')
        ds = AugmentImageComponent(ds, [
            imgaug.ResizeShortestEdge(224, cv2.INTER_CUBIC),
            imgaug.CenterCrop((224, 224)),
        ])
        ds = PrefetchData(ds, 100, multiprocessing.cpu_count())
        ds = BatchData(ds, batchsize)

    return ds
示例#19
0
def get_dataflow_batch(path, is_train, batchsize, img_path=None):
    logger.info('dataflow img_path=%s' % img_path)
    ds = get_dataflow(path, is_train, img_path=img_path)
    print("ds from get_dataflow", ds)
    ds = BatchData(ds, batchsize)
    print("ds from batchdata", ds)
    if is_train:
        ds = PrefetchData(ds, 10, 3)
        print("ds from preferchdata", ds)
    else:
        ds = PrefetchData(ds, 50, 2)

    return ds
示例#20
0
def get_dataflow_batch(path, is_train=True, batch_size=10, img_path=None,sigma=8.0,output_shape=(1440,2560),
                                   numparts=5,translation=False,scale=False,rotation=True,
                                                mins=0.25,maxs=1.2,mina=-np.pi,maxa=np.pi, ilumination=0.0,image_type='RGB'):
    logger.info('dataflow img_path=%s' % img_path)
    
    ds = get_dataflow(path, is_train, img_path=img_path,sigma=sigma,output_shape=output_shape, 
                            translation=translation,scale=scale,rotation=rotation,
                                    mins=mins,maxs=maxs, mina=mina,maxa=maxa,ilumination=ilumination,image_type=image_type)
    ds = BatchData(ds, batch_size)
    # if is_train:
    ds = PrefetchData(ds, 10, 2)
    # else:
    #     ds = PrefetchData(ds, 50, 2)

    return ds
示例#21
0
def get_infer_iterator(hparams, dataset, num_gpu, batch_size):

    df = DataFromList(dataset, shuffle=False)
    num_samples = len(df)
    if num_samples % batch_size != 0 and num_samples % batch_size < num_gpu:
        raise ValueError("num_samples %% batch_size < num_gpu")

    df = MapData(df, lambda data: map_func(hparams, data))
    batched_df = BatchData(df, batch_size=batch_size, remainder=True)
    splitted_df = MapData(
        batched_df,
        lambda x: [np.array_split(x[idx], num_gpu) for idx in range(len(x))])
    prefetched_df = PrefetchDataZMQ(splitted_df,
                                    nr_proc=1,
                                    hwm=batch_size * 10)

    return prefetched_df
示例#22
0
def get_config():
    logger.set_logger_dir(LOG_DIR)
    M = Model()

    name_base = str(uuid.uuid1())[:6]
    PIPE_DIR = os.environ.get('TENSORPACK_PIPEDIR', '.').rstrip('/')
    namec2s = 'ipc://{}/sim-c2s-{}'.format(PIPE_DIR, name_base)
    names2c = 'ipc://{}/sim-s2c-{}'.format(PIPE_DIR, name_base)
    procs = [
        MySimulatorWorker(k, namec2s, names2c) for k in range(SIMULATOR_PROC)
    ]
    ensure_proc_terminate(procs)
    start_proc_mask_signal(procs)

    master = MySimulatorMaster(namec2s, names2c, M)
    dataflow = BatchData(DataFromQueue(master.queue), BATCH_SIZE)

    lr = tf.Variable(0.001, trainable=False, name='learning_rate')
    tf.scalar_summary('learning_rate', lr)

    return TrainConfig(
        dataset=dataflow,
        optimizer=tf.train.AdamOptimizer(lr, epsilon=1e-3),
        callbacks=Callbacks([
            StatPrinter(),
            PeriodicCallback(ModelSaver(), 5),
            ScheduledHyperParamSetter('learning_rate', [(80, 0.0003),
                                                        (120, 0.0001)]),
            ScheduledHyperParamSetter('entropy_beta', [(80, 0.005)]),
            ScheduledHyperParamSetter('explore_factor', [(80, 2), (100, 3),
                                                         (120, 4), (140, 5)]),
            HumanHyperParamSetter('learning_rate'),
            HumanHyperParamSetter('entropy_beta'),
            HumanHyperParamSetter('explore_factor'),
            master,
            PeriodicCallback(
                Evaluator(EVAL_EPISODE, ['state'], ['logits'],
                          policy_dist=POLICY_DIST), 5),
        ]),
        extra_threads_procs=[master],
        session_config=get_default_sess_config(0.5),
        model=M,
        step_per_epoch=STEP_PER_EPOCH,
        max_epoch=1000,
    )
def get_infer_iterator(dataset, hparams, lmdb_path):

    serialize_to_lmdb(dataset, hparams, lmdb_path)

    batch_size = hparams.infer_batch_size
    num_gpu = hparams.num_gpu

    df = LMDBSerializer.load(lmdb_path, shuffle=False)

    batched_df = BatchData(df, batch_size=batch_size, remainder=False)
    splitted_df = MapData(
        batched_df,
        lambda x: [np.array_split(x[idx], num_gpu) for idx in range(len(x))])
    prefetched_df = PrefetchDataZMQ(splitted_df,
                                    nr_proc=1,
                                    hwm=batch_size * 10)

    return prefetched_df
示例#24
0
def batch_dataflow(df,
                   batch_size,
                   time_steps=4,
                   num_stages=6,
                   format=['heatpaf', 'last']):
    informat, outformat = format

    df = BatchData(df, batch_size, use_list=False)

    def in_heat(x):
        return [
            np.stack([x[0]] * time_steps, axis=1),
            np.stack([x[2]] * time_steps, axis=1)
        ]

    def in_heatpaf(x):
        return [
            np.stack([x[0]] * time_steps, axis=1),
            np.stack([x[1]] * time_steps, axis=1),
            np.stack([x[2]] * time_steps, axis=1)
        ]

    def out_heat_last(x):
        return [np.stack([x[4]] * time_steps, axis=1)] * num_stages

    def out_heatpaf_last(x):
        return [
            np.stack([x[3]] * time_steps, axis=1),
            np.stack([x[4]] * time_steps, axis=1),
            np.stack([x[3]] * time_steps, axis=1),
            np.stack([x[4]] * time_steps, axis=1),  # TD layers end here
            x[3],  # TD layers are joined here by LSTM
            x[4],
            x[3],  # these last outputs collapse to one timestep output
            x[4],
            x[3],
            x[4],
            x[3],
            x[4],
        ]

    if informat == 'heat' and outformat == 'last':
        df = MapData(df, lambda x: (heat_only(x), out_heat_last(x)))
    elif informat == 'heatpaf' and outformat == 'last':
        df = MapData(df, lambda x: (in_heatpaf(x), out_heatpaf_last(x)))
    else:
        raise Exception('Unknown format requested: %s' % format)

    df.reset_state()
    return df
示例#25
0
def get_iterator(hparams,
                 dataset,
                 lmdb_path,
                 shuffle=True,
                 drop_remainder=True,
                 nr_proc=4):

    serialize_to_lmdb(hparams, dataset, lmdb_path)

    batch_size = hparams.batch_size
    num_gpu = hparams.num_gpu
    df = LMDBSerializer.load(lmdb_path, shuffle=shuffle)

    batched_df = BatchData(df,
                           batch_size=batch_size,
                           remainder=not drop_remainder)
    splitted_df = MapData(
        batched_df,
        lambda x: [np.array_split(x[idx], num_gpu) for idx in range(len(x))])
    prefetched_df = PrefetchDataZMQ(splitted_df,
                                    nr_proc=nr_proc,
                                    hwm=batch_size * 10)

    return prefetched_df
示例#26
0
 def __call__(self, n_prefetch=1000, n_thread=1):
     df = self
     df = BatchData(df, self.batch_size)
     df = PrefetchData(df, n_prefetch, n_thread)
     return df
示例#27
0
def get_config(args=None,
               is_chief=True,
               task_index=0,
               chief_worker_hostname="",
               n_workers=1):
    logger.set_logger_dir(args.train_log_path +
                          datetime.now().strftime('%Y-%m-%d_%H-%M-%S') + '_' +
                          str(task_index))

    # function to split model parameters between multiple parameter servers
    ps_strategy = tf.contrib.training.GreedyLoadBalancingStrategy(
        len(cluster['ps']), tf.contrib.training.byte_size_load_fn)
    device_function = tf.train.replica_device_setter(
        worker_device='/job:worker/task:{}/cpu:0'.format(task_index),
        cluster=cluster_spec,
        ps_strategy=ps_strategy)

    M = Model(device_function)

    name_base = str(uuid.uuid1()).replace('-', '')[:16]
    PIPE_DIR = os.environ.get('TENSORPACK_PIPEDIR', '.').rstrip('/')
    namec2s = 'ipc://{}/sim-c2s-{}'.format(PIPE_DIR, name_base)
    names2c = 'ipc://{}/sim-s2c-{}'.format(PIPE_DIR, name_base)
    procs = [
        MySimulatorWorker(k, namec2s, names2c)
        for k in range(args.simulator_procs)
    ]
    ensure_proc_terminate(procs)
    start_proc_mask_signal(procs)

    neptune_client = neptune_mp_server.Client(
        server_host=chief_worker_hostname, server_port=args.port)

    master = MySimulatorMaster(task_index,
                               neptune_client,
                               namec2s,
                               names2c,
                               M,
                               dummy=args.dummy,
                               predictor_threads=args.nr_predict_towers,
                               predict_batch_size=args.predict_batch_size,
                               do_train=args.do_train)

    # here's the data passed to the repeated data source
    dataflow = BatchData(DataFromQueue(master.queue), BATCH_SIZE)

    with tf.device(device_function):
        with tf.variable_scope(tf.get_variable_scope(), reuse=None):
            lr = tf.Variable(args.learning_rate,
                             trainable=False,
                             name='learning_rate')
    tf.summary.scalar('learning_rate', lr)

    intra_op_par = args.intra_op_par
    inter_op_par = args.inter_op_par

    session_config = get_default_sess_config(0.5)
    print("{} {}".format(intra_op_par, type(intra_op_par)))
    if intra_op_par is not None:
        session_config.intra_op_parallelism_threads = intra_op_par

    if inter_op_par is not None:
        session_config.inter_op_parallelism_threads = inter_op_par

    session_config.log_device_placement = False
    extra_arg = {
        'dummy_predictor': args.dummy_predictor,
        'intra_op_par': intra_op_par,
        'inter_op_par': inter_op_par,
        'max_steps': args.max_steps,
        'device_count': {
            'CPU': args.cpu_device_count
        },
        'threads_to_trace': args.threads_to_trace,
        'dummy': args.dummy,
        'cpu': args.cpu,
        'queue_size': args.queue_size,
        #'worker_host' : "grpc://localhost:{}".format(cluster['worker'][my_task_index].split(':')[1]),
        'worker_host': server.target,
        'is_chief': is_chief,
        'device_function': device_function,
        'n_workers': n_workers,
        'use_sync_opt': args.use_sync_opt,
        'port': args.port,
        'batch_size': BATCH_SIZE,
        'debug_charts': args.debug_charts,
        'adam_debug': args.adam_debug,
        'task_index': task_index,
        'lr': lr,
        'schedule_hyper': args.schedule_hyper,
        'experiment_dir': args.experiment_dir
    }

    print("\n\n worker host: {} \n\n".format(extra_arg['worker_host']))

    with tf.device(device_function):
        if args.optimizer == 'adam':
            optimizer = tf.train.AdamOptimizer(lr,
                                               epsilon=args.epsilon,
                                               beta1=args.beta1,
                                               beta2=args.beta2)
            if args.adam_debug:
                optimizer = MyAdamOptimizer(lr,
                                            epsilon=args.epsilon,
                                            beta1=args.beta1,
                                            beta2=args.beta2)
        elif args.optimizer == 'gd':
            optimizer = tf.train.GradientDescentOptimizer(lr)
        elif args.optimizer == 'adagrad':
            optimizer = tf.train.AdagradOptimizer(lr)
        elif args.optimizer == 'adadelta':
            optimizer = tf.train.AdadeltaOptimizer(lr, epsilon=1e-3)
        elif args.optimizer == 'momentum':
            optimizer = tf.train.MomentumOptimizer(lr, momentum=0.9)
        elif args.optimizer == 'rms':
            optimizer = tf.train.RMSPropOptimizer(lr)

        # wrap in SyncReplicasOptimizer
        if args.use_sync_opt == 1:
            if not args.adam_debug:
                optimizer = tf.train.SyncReplicasOptimizer(
                    optimizer,
                    replicas_to_aggregate=args.num_grad,
                    total_num_replicas=n_workers)
            else:
                optimizer = MySyncReplicasOptimizer(
                    optimizer,
                    replicas_to_aggregate=args.num_grad,
                    total_num_replicas=n_workers)
            extra_arg['hooks'] = optimizer.make_session_run_hook(is_chief)

    callbacks = [
        StatPrinter(), master,
        DebugLogCallback(neptune_client,
                         worker_id=task_index,
                         nr_send=args.send_debug_every,
                         debug_charts=args.debug_charts,
                         adam_debug=args.adam_debug,
                         schedule_hyper=args.schedule_hyper)
    ]

    if args.debug_charts:
        callbacks.append(
            HeartPulseCallback('heart_pulse_{}.log'.format(
                os.environ['SLURMD_NODENAME'])))

    if args.early_stopping is not None:
        args.early_stopping = float(args.early_stopping)

        if my_task_index == 1 and not args.eval_node:
            # only one worker does evaluation
            callbacks.append(
                PeriodicCallback(
                    Evaluator(EVAL_EPISODE, ['state'], ['logits'],
                              neptune_client,
                              worker_id=task_index,
                              solved_score=args.early_stopping), 2))
    elif my_task_index == 1 and not args.eval_node:
        # only 1 worker does evaluation
        callbacks.append(
            PeriodicCallback(
                Evaluator(EVAL_EPISODE, ['state'], ['logits'],
                          neptune_client,
                          worker_id=task_index), 2))

    if args.save_every != 0:
        callbacks.append(
            PeriodicPerStepCallback(
                ModelSaver(var_collections=M.vars_for_save,
                           models_dir=args.models_dir), args.save_every))

    if args.schedule_hyper and my_task_index == 2:
        callbacks.append(
            HyperParameterScheduler('learning_rate', [(20, 0.0005),
                                                      (60, 0.0001)]))
        callbacks.append(
            HyperParameterScheduler('entropy_beta', [(40, 0.005),
                                                     (80, 0.001)]))

    return TrainConfig(dataset=dataflow,
                       optimizer=optimizer,
                       callbacks=Callbacks(callbacks),
                       extra_threads_procs=[master],
                       session_config=session_config,
                       model=M,
                       step_per_epoch=STEP_PER_EPOCH,
                       max_epoch=args.max_epoch,
                       extra_arg=extra_arg)
示例#28
0
 def __call__(self, n_prefetch=1, n_thread=1):
     df = self
     df = BatchData(df, 1)
     df = PrefetchData(df, n_prefetch, n_thread)
     return df
示例#29
0
def get_dataflow_batch(path, is_train, batchsize):
    ds = get_dataflow(path, is_train)
    ds = BatchData(ds, batchsize)
    ds = PrefetchData(ds, 10, 2)

    return ds
示例#30
0
 def dataflow(self, nr_prefetch=1000, nr_thread=1):
     ds = self
     ds = BatchData(ds, self.batch_size)
     ds = PrefetchData(ds, nr_prefetch, nr_thread)
     return ds