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
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
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
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
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
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
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
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
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
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
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
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
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
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
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
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
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
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
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
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
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
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)
def __call__(self, n_prefetch=1, n_thread=1):
df = self
df = BatchData(df, 1)
df = PrefetchData(df, n_prefetch, n_thread)
return df
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
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
