def get_estimator(batch_size=4, epochs=1000, steps_per_epoch=1000, validation_steps=None, model_dir=tempfile.mkdtemp(), imagenet_path=None):
"""Args:
imagenet_path: folder path of ImageNet dataset, containing train and val subdirs .
"""
assert imagenet_path is not None, 'Pass valid folder path of Imagenet dataset'
# Ensure ImageNet dataset is downloaded. Pass the folder contianing train and val subdirectories.
# currently the script doesn't download the ImageNet data.
train_csv, val_csv, path = srgan.load_data(path_imgnet= imagenet_path)
writer = fe.RecordWriter(
save_dir=os.path.join(path, "sr_tfrecords"),
train_data=train_csv,
validation_data=val_csv,
ops=[ImageReader(inputs="lowres", outputs="lowres"), ImageReader(inputs="highres", outputs="highres")],
compression="GZIP",
write_feature=['lowres', 'highres'])
pipeline = fe.Pipeline(
max_shuffle_buffer_mb=3000,
batch_size=batch_size,
data=writer,
ops=[
LowresRescale(inputs='lowres', outputs='lowres'),
Rescale(inputs='highres', outputs='highres'), ])
# prepare model
model = fe.build(model_def=lambda: get_generator(input_shape=(24, 24, 3)),
model_name="srresnet_gen",
optimizer=tf.optimizers.Adam(learning_rate=0.0001),
loss_name="mse_loss")
network = fe.Network(ops=[
ModelOp(inputs='lowres', model=model, outputs='superres'),
PixelMeanSquaredError(inputs=('superres', 'highres'), outputs="mse_loss")
])
model_dir = os.path.join(path)
estimator = fe.Estimator(network=network,
pipeline=pipeline,
steps_per_epoch=steps_per_epoch,
epochs=epochs,
traces=[
ModelSaver(model_name="srresnet_gen", save_dir=model_dir, save_best=True), ])
return estimator
def get_estimator(batch_size=4,
epochs=200,
steps_per_epoch=1000,
validation_steps=None,
model_dir=tempfile.mkdtemp(),
imagenet_path=None,
srresnet_model_path=None):
"""Args:
imagenet_path: folder path of ImageNet dataset, containing train and val subdirs .
srresnet_model_path: srresnet model weights, srgan generator gets initialized with the weights.
"""
assert imagenet_path is not None, 'Pass valid folder path of Imagenet dataset'
assert srresnet_model_path is not None, 'srresnet model is needed to initialize srgan generator model'
# Ensure ImageNet dataset is downloaded. Pass the folder contianing train and val subdirectories.
# currently the script doesn't download the ImageNet data.
train_csv, val_csv, path = srgan.load_data(path_imgnet=imagenet_path)
writer = fe.RecordWriter(save_dir=os.path.join(path, "sr_tfrecords"),
train_data=train_csv,
validation_data=val_csv,
ops=[
ImageReader(inputs="lowres",
outputs="lowres"),
ImageReader(inputs="highres",
outputs="highres")
],
compression="GZIP",
write_feature=['lowres', 'highres'])
pipeline = fe.Pipeline(max_shuffle_buffer_mb=3000,
batch_size=batch_size,
data=writer,
ops=[
LowresRescale(inputs='lowres',
outputs='lowres'),
Rescale(inputs='highres', outputs='highres'),
])
# prepare model
model_gen = fe.build(model_def=srresnet_model_path,
model_name="srgan_gen",
optimizer=tf.optimizers.Adam(learning_rate=0.0001),
loss_name="mse_adv_loss",
custom_objects={'SubPixelConv2D': SubPixelConv2D})
model_desc = fe.build(
model_def=lambda: get_discriminator(input_shape=(96, 96, 3)),
model_name="srgan_desc",
optimizer=tf.optimizers.Adam(learning_rate=0.0001),
loss_name="desc_loss")
network = fe.Network(ops=[
ModelOp(inputs='lowres', model=model_gen, outputs='superres'),
ModelOp(inputs='superres', model=model_desc, outputs='pred_fake'),
ModelOp(inputs='highres', model=model_desc, outputs='pred_true'),
DLoss(inputs=("pred_true", "pred_fake"),
outputs=("desc_loss", "real_loss", "fake_loss")),
GLoss(inputs=('superres', 'highres', 'pred_fake'),
outputs=("mse_adv_loss", "mse_loss", "adv_loss"),
vgg_content=True)
])
model_dir = os.path.join(path)
estimator = fe.Estimator(
network=network,
pipeline=pipeline,
steps_per_epoch=steps_per_epoch,
epochs=epochs,
traces=[
ModelSaver(model_name="srgan_gen",
save_dir=model_dir,
save_best=True),
ModelSaver(model_name="srgan_desc",
save_dir=model_dir,
save_best=True),
LRController(model_name="srgan_gen",
lr_schedule=MyLRSchedule(schedule_mode='step')),
LRController(model_name="srgan_desc",
lr_schedule=MyLRSchedule(schedule_mode='step'))
])
return estimator
def get_estimator(batch_size=1,
epochs=500,
steps_per_epoch=128,
model_dir=tempfile.mkdtemp(),
path_brats=os.path.join(os.getenv('HOME'),
'fastestimator_data', 'BraTs')):
"""Args:
path_brats: folder path of BraTs 2018 dataset, containing data subdir inturn having LGG and HGG data.
Expected folder structure path_brats
path_brats/
|----------data/
|----LGG/
|----HGG/
"""
assert path_brats is not None, 'Pass valid folder path of BraTs 2018 dataset'
# Ensure Brats 2018 dataset is downloaded. Pass the folder contianing train and val subdirectories.
# currently the script doesn't download the BraTs data.
train_csv, val_csv, path_brats = brats.load_data(path_brats=path_brats,
resized_img_shape=(128,
128,
128),
bias_correction=False)
writer = fe.RecordWriter(save_dir=os.path.join(path_brats, "tfrecords"),
train_data=train_csv,
validation_data=val_csv,
ops=[
NIBImageReader(inputs="mod_img",
outputs="mod_img"),
NIBImageReader(inputs="seg_mask",
outputs="seg_mask")
],
compression="GZIP",
write_feature=['mod_img', 'seg_mask'])
pipeline = fe.Pipeline(data=writer,
batch_size=batch_size,
ops=[
SplitMaskLabelwise(inputs="seg_mask",
outputs="seg_mask"),
Augmentation3D(inputs=("mod_img", "seg_mask"),
outputs=("mod_img", "seg_mask"),
mode='train')
])
model_unet3d_is = fe.build(
model_def=lambda: UNet3D_Isensee(input_shape=(4, 64, 64, 64)),
model_name="brats_unet3d_is",
optimizer=tf.optimizers.Adam(learning_rate=5e-4),
loss_name="wdice_loss")
network = fe.Network(ops=[
ModelOp(
inputs="mod_img", model=model_unet3d_is, outputs="pred_seg_mask"),
WeightedDiceLoss(inputs=("seg_mask", "pred_seg_mask"),
outputs=("wdice_loss"))
])
model_dir = path_brats
estimator = fe.Estimator(network=network,
pipeline=pipeline,
steps_per_epoch=steps_per_epoch,
epochs=epochs,
traces=[
ModelSaver(model_name="brats_unet3d_is",
save_dir=model_dir,
save_best=True),
LRController(model_name="brats_unet3d_is",
reduce_patience=10,
reduce_factor=0.5,
reduce_on_eval=True,
min_lr=1e-07)
],
log_steps=steps_per_epoch)
return estimator
def get_estimator(data_path=None, model_dir=tempfile.mkdtemp(), batch_size=2):
#prepare dataset
train_csv, val_csv, path = load_data(path=data_path)
writer = fe.RecordWriter(
save_dir=os.path.join(path, "retinanet_coco_1024"),
train_data=train_csv,
validation_data=val_csv,
ops=[
ImageReader(inputs="image", parent_path=path, outputs="image"),
String2List(inputs=["x1", "y1", "width", "height", "obj_label"],
outputs=["x1", "y1", "width", "height", "obj_label"]),
ResizeImageAndBbox(
target_size=(1024, 1024),
keep_ratio=True,
inputs=["image", "x1", "y1", "width", "height"],
outputs=["image", "x1", "y1", "width", "height"]),
FlipImageAndBbox(inputs=[
"image", "x1", "y1", "width", "height", "obj_label", "id"
],
outputs=[
"image", "x1", "y1", "width", "height",
"obj_label", "id"
]),
GenerateTarget(inputs=("obj_label", "x1", "y1", "width", "height"),
outputs=("cls_gt", "x1_gt", "y1_gt", "w_gt",
"h_gt"))
],
expand_dims=True,
compression="GZIP",
write_feature=[
"image", "id", "cls_gt", "x1_gt", "y1_gt", "w_gt", "h_gt",
"obj_label", "x1", "y1", "width", "height"
])
# prepare pipeline
pipeline = fe.Pipeline(batch_size=batch_size,
data=writer,
ops=[
Rescale(inputs="image", outputs="image"),
Pad(padded_shape=[2051],
inputs=[
"x1_gt", "y1_gt", "w_gt", "h_gt",
"obj_label", "x1", "y1", "width",
"height"
],
outputs=[
"x1_gt", "y1_gt", "w_gt", "h_gt",
"obj_label", "x1", "y1", "width",
"height"
])
])
# prepare network
model = fe.build(model_def=lambda: RetinaNet(input_shape=(1024, 1024, 3),
num_classes=90),
model_name="retinanet",
optimizer=tf.optimizers.SGD(momentum=0.9),
loss_name="total_loss")
network = fe.Network(ops=[
ModelOp(inputs="image", model=model, outputs=["cls_pred", "loc_pred"]),
RetinaLoss(inputs=("cls_gt", "x1_gt", "y1_gt", "w_gt", "h_gt",
"cls_pred", "loc_pred"),
outputs=("total_loss", "focal_loss", "l1_loss")),
PredictBox(inputs=[
"cls_pred", "loc_pred", "obj_label", "x1", "y1", "width", "height"
],
outputs=("pred", "gt"),
mode="eval",
input_shape=(1024, 1024, 3))
])
# prepare estimator
estimator = fe.Estimator(
network=network,
pipeline=pipeline,
epochs=7,
traces=[
MeanAvgPrecision(90, (1024, 1024, 3),
'pred',
'gt',
output_name=("mAP", "AP50", "AP75")),
ModelSaver(model_name="retinanet",
save_dir=model_dir,
save_best='mAP',
save_best_mode='max'),
LRController(model_name="retinanet",
lr_schedule=MyLRSchedule(schedule_mode="step"))
])
return estimator
def get_estimator(batch_size=128, epochs=100):
usps_train_csv, _, usps_parent_dir = usps.load_data()
mnist_train_csv, _, mnist_parent_dir = mnist.load_data()
df = pd.read_csv(mnist_train_csv)
df.columns = ['source_img', 'source_label']
df.to_csv(mnist_train_csv, index=False)
df = pd.read_csv(usps_train_csv)
df.columns = ['target_img', 'target_label']
df.to_csv(usps_train_csv, index=False)
writer = fe.RecordWriter(
save_dir=os.path.join(os.path.dirname(mnist_parent_dir), 'dann',
'tfr'),
train_data=(usps_train_csv, mnist_train_csv),
ops=(
[
ImageReader(inputs="target_img",
outputs="target_img",
parent_path=usps_parent_dir,
grey_scale=True)
], # first tuple element
[
ImageReader(inputs="source_img",
outputs="source_img",
parent_path=mnist_parent_dir,
grey_scale=True)
])) # second tuple element
pipeline = fe.Pipeline(batch_size=batch_size,
data=writer,
ops=[
Resize(inputs="target_img",
outputs="target_img",
size=(28, 28)),
Resize(inputs="source_img",
outputs="source_img",
size=(28, 28)),
Minmax(inputs="target_img",
outputs="target_img"),
Minmax(inputs="source_img",
outputs="source_img")
])
alpha = tf.Variable(0.0, dtype=tf.float32, trainable=False)
img_shape = (28, 28, 1)
feat_dim = 7 * 7 * 48
feature_extractor = fe.build(
model_def=lambda: build_feature_extractor(img_shape),
model_name="feature_extractor",
loss_name="fe_loss",
optimizer=tf.keras.optimizers.Adam(1e-4))
label_predictor = fe.build(
model_def=lambda: build_label_predictor(feat_dim),
model_name="label_predictor",
loss_name="fe_loss",
optimizer=tf.keras.optimizers.Adam(1e-4))
domain_predictor = fe.build(
model_def=lambda: build_domain_predictor(feat_dim, alpha),
model_name="domain_predictor",
loss_name="fe_loss",
optimizer=tf.keras.optimizers.Adam(1e-4))
network = fe.Network(ops=[
ModelOp(
inputs="source_img", outputs="src_feat", model=feature_extractor),
ModelOp(
inputs="target_img", outputs="tgt_feat", model=feature_extractor),
ModelOp(
inputs="src_feat", outputs="src_c_logit", model=label_predictor),
ModelOp(
inputs="src_feat", outputs="src_d_logit", model=domain_predictor),
ModelOp(
inputs="tgt_feat", outputs="tgt_d_logit", model=domain_predictor),
FELoss(inputs=("src_c_logit", "source_label", "src_d_logit",
"tgt_d_logit"),
outputs="fe_loss")
])
traces = [GRLWeightController(alpha=alpha)]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
traces=traces,
epochs=epochs)
return estimator