def test_input(self):
resize = Resize(image_in='x', height=self.height, width=self.width)
output = resize.forward(data=self.single_input, state={})
with self.subTest('Check output type'):
self.assertEqual(type(output), list)
with self.subTest('Check output image shape'):
self.assertEqual(output[0].shape, self.single_output_shape)
def test_input_image_and_mask(self):
resize = Resize(image_in='x', mask_in='x_mask', height=self.height, width=self.width)
output = resize.forward(data=self.input_image_and_mask, state={})
with self.subTest('Check output type'):
self.assertEqual(type(output), list)
with self.subTest('Check output image shape'):
self.assertEqual(output[0].shape, self.image_and_mask_output_shape)
with self.subTest('Check output mask shape'):
self.assertEqual(output[1].shape, self.image_and_mask_output_shape)
def get_estimator(batch_size=4,
epochs=2,
max_train_steps_per_epoch=None,
log_steps=100,
style_weight=5.0,
content_weight=1.0,
tv_weight=1e-4,
save_dir=tempfile.mkdtemp(),
style_img_path='Vassily_Kandinsky,_1913_-_Composition_7.jpg',
data_dir=None):
train_data, _ = mscoco.load_data(root_dir=data_dir, load_bboxes=False, load_masks=False, load_captions=False)
device = "cuda" if torch.cuda.is_available() else "cpu"
style_img = cv2.imread(style_img_path)
assert style_img is not None, "cannot load the style image, please go to the folder with style image"
style_img = cv2.resize(style_img, (256, 256))
style_img = (style_img.astype(np.float32) - 127.5) / 127.5
pipeline = fe.Pipeline(
train_data=train_data,
batch_size=batch_size,
ops=[
ReadImage(inputs="image", outputs="image"),
Normalize(inputs="image", outputs="image", mean=1.0, std=1.0, max_pixel_value=127.5),
Resize(height=256, width=256, image_in="image", image_out="image"),
LambdaOp(fn=lambda: style_img, outputs="style_image"),
ChannelTranspose(inputs=["image", "style_image"], outputs=["image", "style_image"])
])
model = fe.build(model_fn=StyleTransferNet,
model_name="style_transfer_net",
optimizer_fn=lambda x: torch.optim.Adam(x, lr=1e-3))
network = fe.Network(ops=[
ModelOp(inputs="image", model=model, outputs="image_out"),
ExtractVGGFeatures(inputs="style_image", outputs="y_style", device=device),
ExtractVGGFeatures(inputs="image", outputs="y_content", device=device),
ExtractVGGFeatures(inputs="image_out", outputs="y_pred", device=device),
StyleContentLoss(style_weight=style_weight,
content_weight=content_weight,
tv_weight=tv_weight,
inputs=('y_pred', 'y_style', 'y_content', 'image_out'),
outputs='loss'),
UpdateOp(model=model, loss_name="loss")
])
estimator = fe.Estimator(network=network,
pipeline=pipeline,
traces=ModelSaver(model=model, save_dir=save_dir, frequency=1),
epochs=epochs,
max_train_steps_per_epoch=max_train_steps_per_epoch,
log_steps=log_steps)
return estimator
def get_estimator(epochs=20,
batch_size=4,
train_steps_per_epoch=None,
eval_steps_per_epoch=None,
save_dir=tempfile.mkdtemp(),
log_steps=20,
data_dir=None):
# step 1
csv = montgomery.load_data(root_dir=data_dir)
pipeline = fe.Pipeline(
train_data=csv,
eval_data=csv.split(0.2),
batch_size=batch_size,
ops=[
ReadImage(inputs="image",
parent_path=csv.parent_path,
outputs="image",
color_flag='gray'),
ReadImage(inputs="mask_left",
parent_path=csv.parent_path,
outputs="mask_left",
color_flag='gray',
mode='!infer'),
ReadImage(inputs="mask_right",
parent_path=csv.parent_path,
outputs="mask_right",
color_flag='gray',
mode='!infer'),
CombineLeftRightMask(inputs=("mask_left", "mask_right"),
outputs="mask",
mode='!infer'),
Resize(image_in="image", width=512, height=512),
Resize(image_in="mask", width=512, height=512, mode='!infer'),
Sometimes(numpy_op=HorizontalFlip(
image_in="image", mask_in="mask", mode='train')),
Sometimes(numpy_op=Rotate(image_in="image",
mask_in="mask",
limit=(-10, 10),
border_mode=cv2.BORDER_CONSTANT,
mode='train')),
Minmax(inputs="image", outputs="image"),
Minmax(inputs="mask", outputs="mask", mode='!infer')
])
# step 2
model = fe.build(
model_fn=lambda: UNet(input_size=(512, 512, 1)),
optimizer_fn=lambda: tf.keras.optimizers.Adam(learning_rate=0.0001),
model_name="lung_segmentation")
network = fe.Network(ops=[
ModelOp(inputs="image", model=model, outputs="pred_segment"),
CrossEntropy(
inputs=("pred_segment", "mask"), outputs="loss", form="binary"),
UpdateOp(model=model, loss_name="loss")
])
# step 3
traces = [
Dice(true_key="mask", pred_key="pred_segment"),
BestModelSaver(model=model,
save_dir=save_dir,
metric='Dice',
save_best_mode='max')
]
estimator = fe.Estimator(network=network,
pipeline=pipeline,
epochs=epochs,
log_steps=log_steps,
traces=traces,
train_steps_per_epoch=train_steps_per_epoch,
eval_steps_per_epoch=eval_steps_per_epoch)
return estimator
def get_estimator(target_size=128,
epochs=55,
save_dir=tempfile.mkdtemp(),
max_train_steps_per_epoch=None,
data_dir=None):
# assert growth parameters
num_grow = np.log2(target_size) - 2
assert num_grow >= 1 and num_grow % 1 == 0, "need exponential of 2 and greater than 8 as target size"
num_phases = int(2 * num_grow + 1)
assert epochs % num_phases == 0, "epoch must be multiple of {} for size {}".format(
num_phases, target_size)
num_grow, phase_length = int(num_grow), int(epochs / num_phases)
event_epoch = [1, 1 + phase_length] + [
phase_length * (2 * i + 1) + 1 for i in range(1, num_grow)
]
event_size = [4] + [2**(i + 3) for i in range(num_grow)]
# set up data schedules
dataset = nih_chestxray.load_data(root_dir=data_dir)
resize_map = {
epoch: Resize(image_in="x", image_out="x", height=size, width=size)
for (epoch, size) in zip(event_epoch, event_size)
}
resize_low_res_map1 = {
epoch: Resize(image_in="x",
image_out="x_low_res",
height=size // 2,
width=size // 2)
for (epoch, size) in zip(event_epoch, event_size)
}
resize_low_res_map2 = {
epoch: Resize(image_in="x_low_res",
image_out="x_low_res",
height=size,
width=size)
for (epoch, size) in zip(event_epoch, event_size)
}
batch_size_map = {
epoch: 512 // size * get_num_devices() if size <= 128 else 4 *
get_num_devices()
for (epoch, size) in zip(event_epoch, event_size)
}
batch_scheduler = EpochScheduler(epoch_dict=batch_size_map)
pipeline = fe.Pipeline(
batch_size=batch_scheduler,
train_data=dataset,
drop_last=True,
ops=[
ReadImage(inputs="x", outputs="x", color_flag='gray'),
EpochScheduler(epoch_dict=resize_map),
EpochScheduler(epoch_dict=resize_low_res_map1),
EpochScheduler(epoch_dict=resize_low_res_map2),
Normalize(inputs=["x", "x_low_res"],
outputs=["x", "x_low_res"],
mean=1.0,
std=1.0,
max_pixel_value=127.5),
LambdaOp(fn=lambda: np.random.normal(size=[512]).astype('float32'),
outputs="z")
])
# now model schedule
fade_in_alpha = tf.Variable(initial_value=1.0,
dtype='float32',
trainable=False)
d_models = fe.build(
model_fn=lambda: build_D(fade_in_alpha,
target_resolution=int(np.log2(target_size)),
num_channels=1),
optimizer_fn=[
lambda: Adam(0.001, beta_1=0.0, beta_2=0.99, epsilon=1e-8)
] * len(event_size),
model_name=["d_{}".format(size) for size in event_size])
g_models = fe.build(
model_fn=lambda: build_G(fade_in_alpha,
target_resolution=int(np.log2(target_size)),
num_channels=1),
optimizer_fn=[
lambda: Adam(0.001, beta_1=0.0, beta_2=0.99, epsilon=1e-8)
] * len(event_size) + [None],
model_name=["g_{}".format(size) for size in event_size] + ["G"])
fake_img_map = {
epoch: ModelOp(inputs="z", outputs="x_fake", model=model)
for (epoch, model) in zip(event_epoch, g_models[:-1])
}
fake_score_map = {
epoch: ModelOp(inputs="x_fake", outputs="fake_score", model=model)
for (epoch, model) in zip(event_epoch, d_models)
}
real_score_map = {
epoch: ModelOp(inputs="x_blend", outputs="real_score", model=model)
for (epoch, model) in zip(event_epoch, d_models)
}
interp_score_map = {
epoch: ModelOp(inputs="x_interp", outputs="interp_score", model=model)
for (epoch, model) in zip(event_epoch, d_models)
}
g_update_map = {
epoch: UpdateOp(loss_name="gloss", model=model)
for (epoch, model) in zip(event_epoch, g_models[:-1])
}
d_update_map = {
epoch: UpdateOp(loss_name="dloss", model=model)
for (epoch, model) in zip(event_epoch, d_models)
}
network = fe.Network(ops=[
EpochScheduler(fake_img_map),
EpochScheduler(fake_score_map),
ImageBlender(
alpha=fade_in_alpha, inputs=("x", "x_low_res"), outputs="x_blend"),
EpochScheduler(real_score_map),
Interpolate(inputs=("x_fake", "x"), outputs="x_interp"),
EpochScheduler(interp_score_map),
GradientPenalty(inputs=("x_interp", "interp_score"), outputs="gp"),
GLoss(inputs="fake_score", outputs="gloss"),
DLoss(inputs=("real_score", "fake_score", "gp"), outputs="dloss"),
EpochScheduler(g_update_map),
EpochScheduler(d_update_map)
])
traces = [
AlphaController(alpha=fade_in_alpha,
fade_start_epochs=event_epoch[1:],
duration=phase_length,
batch_scheduler=batch_scheduler,
num_examples=len(dataset)),
ModelSaver(model=g_models[-1],
save_dir=save_dir,
frequency=phase_length),
ImageSaving(epoch_model_map={
epoch - 1: model
for (epoch,
model) in zip(event_epoch[1:] + [epochs + 1], g_models[:-1])
},
save_dir=save_dir)
]
estimator = fe.Estimator(
pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch)
return estimator
def get_estimator(weight=10.0,
epochs=200,
batch_size=1,
max_train_steps_per_epoch=None,
save_dir=tempfile.mkdtemp(),
data_dir=None):
train_data, _ = load_data(batch_size=batch_size, root_dir=data_dir)
device = "cuda" if torch.cuda.is_available() else "cpu"
pipeline = fe.Pipeline(train_data=train_data,
ops=[
ReadImage(inputs=["A", "B"], outputs=["A",
"B"]),
Normalize(inputs=["A", "B"],
outputs=["real_A", "real_B"],
mean=1.0,
std=1.0,
max_pixel_value=127.5),
Resize(height=286,
width=286,
image_in="real_A",
image_out="real_A",
mode="train"),
RandomCrop(height=256,
width=256,
image_in="real_A",
image_out="real_A",
mode="train"),
Resize(height=286,
width=286,
image_in="real_B",
image_out="real_B",
mode="train"),
RandomCrop(height=256,
width=256,
image_in="real_B",
image_out="real_B",
mode="train"),
Sometimes(
HorizontalFlip(image_in="real_A",
image_out="real_A",
mode="train")),
Sometimes(
HorizontalFlip(image_in="real_B",
image_out="real_B",
mode="train")),
ChannelTranspose(inputs=["real_A", "real_B"],
outputs=["real_A", "real_B"])
])
g_AtoB = fe.build(model_fn=Generator,
model_name="g_AtoB",
optimizer_fn=lambda x: torch.optim.Adam(
x, lr=2e-4, betas=(0.5, 0.999)))
g_BtoA = fe.build(model_fn=Generator,
model_name="g_BtoA",
optimizer_fn=lambda x: torch.optim.Adam(
x, lr=2e-4, betas=(0.5, 0.999)))
d_A = fe.build(model_fn=Discriminator,
model_name="d_A",
optimizer_fn=lambda x: torch.optim.Adam(
x, lr=2e-4, betas=(0.5, 0.999)))
d_B = fe.build(model_fn=Discriminator,
model_name="d_B",
optimizer_fn=lambda x: torch.optim.Adam(
x, lr=2e-4, betas=(0.5, 0.999)))
network = fe.Network(ops=[
ModelOp(inputs="real_A", model=g_AtoB, outputs="fake_B"),
ModelOp(inputs="real_B", model=g_BtoA, outputs="fake_A"),
Buffer(image_in="fake_A", image_out="buffer_fake_A"),
Buffer(image_in="fake_B", image_out="buffer_fake_B"),
ModelOp(inputs="real_A", model=d_A, outputs="d_real_A"),
ModelOp(inputs="fake_A", model=d_A, outputs="d_fake_A"),
ModelOp(inputs="buffer_fake_A", model=d_A, outputs="buffer_d_fake_A"),
ModelOp(inputs="real_B", model=d_B, outputs="d_real_B"),
ModelOp(inputs="fake_B", model=d_B, outputs="d_fake_B"),
ModelOp(inputs="buffer_fake_B", model=d_B, outputs="buffer_d_fake_B"),
ModelOp(inputs="real_A", model=g_BtoA, outputs="same_A"),
ModelOp(inputs="fake_B", model=g_BtoA, outputs="cycled_A"),
ModelOp(inputs="real_B", model=g_AtoB, outputs="same_B"),
ModelOp(inputs="fake_A", model=g_AtoB, outputs="cycled_B"),
GLoss(inputs=("real_A", "d_fake_B", "cycled_A", "same_A"),
weight=weight,
device=device,
outputs="g_AtoB_loss"),
GLoss(inputs=("real_B", "d_fake_A", "cycled_B", "same_B"),
weight=weight,
device=device,
outputs="g_BtoA_loss"),
DLoss(inputs=("d_real_A", "buffer_d_fake_A"),
outputs="d_A_loss",
device=device),
DLoss(inputs=("d_real_B", "buffer_d_fake_B"),
outputs="d_B_loss",
device=device),
UpdateOp(model=g_AtoB, loss_name="g_AtoB_loss"),
UpdateOp(model=g_BtoA, loss_name="g_BtoA_loss"),
UpdateOp(model=d_A, loss_name="d_A_loss"),
UpdateOp(model=d_B, loss_name="d_B_loss")
])
traces = [
ModelSaver(model=g_AtoB, save_dir=save_dir, frequency=10),
ModelSaver(model=g_BtoA, save_dir=save_dir, frequency=10),
LRScheduler(model=g_AtoB, lr_fn=lr_schedule),
LRScheduler(model=g_BtoA, lr_fn=lr_schedule),
LRScheduler(model=d_A, lr_fn=lr_schedule),
LRScheduler(model=d_B, lr_fn=lr_schedule)
]
estimator = fe.Estimator(
network=network,
pipeline=pipeline,
epochs=epochs,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch)
return estimator
def get_estimator(weight=10.0,
epochs=200,
batch_size=1,
train_steps_per_epoch=None,
save_dir=tempfile.mkdtemp(),
data_dir=None):
train_data, _ = load_data(batch_size=batch_size, root_dir=data_dir)
pipeline = fe.Pipeline(
train_data=train_data,
ops=[
ReadImage(inputs=["A", "B"], outputs=["A", "B"]),
Normalize(inputs=["A", "B"], outputs=["real_A", "real_B"], mean=1.0, std=1.0, max_pixel_value=127.5),
Resize(height=286, width=286, image_in="real_A", image_out="real_A", mode="train"),
RandomCrop(height=256, width=256, image_in="real_A", image_out="real_A", mode="train"),
Resize(height=286, width=286, image_in="real_B", image_out="real_B", mode="train"),
RandomCrop(height=256, width=256, image_in="real_B", image_out="real_B", mode="train"),
Sometimes(HorizontalFlip(image_in="real_A", image_out="real_A", mode="train")),
Sometimes(HorizontalFlip(image_in="real_B", image_out="real_B", mode="train")),
PlaceholderOp(outputs=("index_A", "buffer_A")),
PlaceholderOp(outputs=("index_B", "buffer_B"))
])
g_AtoB = fe.build(model_fn=build_generator, model_name="g_AtoB", optimizer_fn=lambda: tf.optimizers.Adam(2e-4, 0.5))
g_BtoA = fe.build(model_fn=build_generator, model_name="g_BtoA", optimizer_fn=lambda: tf.optimizers.Adam(2e-4, 0.5))
d_A = fe.build(model_fn=build_discriminator, model_name="d_A", optimizer_fn=lambda: tf.optimizers.Adam(2e-4, 0.5))
d_B = fe.build(model_fn=build_discriminator, model_name="d_B", optimizer_fn=lambda: tf.optimizers.Adam(2e-4, 0.5))
network = fe.Network(ops=[
ModelOp(inputs="real_A", model=g_AtoB, outputs="fake_B"),
ModelOp(inputs="real_B", model=g_BtoA, outputs="fake_A"),
Buffer(image_in="fake_A", buffer_in="buffer_A", index_in="index_A", image_out="buffer_fake_A"),
Buffer(image_in="fake_B", buffer_in="buffer_B", index_in="index_B", image_out="buffer_fake_B"),
ModelOp(inputs="real_A", model=d_A, outputs="d_real_A"),
ModelOp(inputs="fake_A", model=d_A, outputs="d_fake_A"),
ModelOp(inputs="buffer_fake_A", model=d_A, outputs="buffer_d_fake_A"),
ModelOp(inputs="real_B", model=d_B, outputs="d_real_B"),
ModelOp(inputs="fake_B", model=d_B, outputs="d_fake_B"),
ModelOp(inputs="buffer_fake_B", model=d_B, outputs="buffer_d_fake_B"),
ModelOp(inputs="real_A", model=g_BtoA, outputs="same_A"),
ModelOp(inputs="fake_B", model=g_BtoA, outputs="cycled_A"),
ModelOp(inputs="real_B", model=g_AtoB, outputs="same_B"),
ModelOp(inputs="fake_A", model=g_AtoB, outputs="cycled_B"),
GLoss(inputs=("real_A", "d_fake_B", "cycled_A", "same_A"), weight=weight, outputs="g_AtoB_loss"),
GLoss(inputs=("real_B", "d_fake_A", "cycled_B", "same_B"), weight=weight, outputs="g_BtoA_loss"),
DLoss(inputs=("d_real_A", "buffer_d_fake_A"), outputs="d_A_loss"),
DLoss(inputs=("d_real_B", "buffer_d_fake_B"), outputs="d_B_loss"),
UpdateOp(model=g_AtoB, loss_name="g_AtoB_loss"),
UpdateOp(model=g_BtoA, loss_name="g_BtoA_loss"),
UpdateOp(model=d_A, loss_name="d_A_loss"),
UpdateOp(model=d_B, loss_name="d_B_loss")
])
traces = [
BufferUpdate(input_name="fake_A",
buffer_size=50,
batch_size=batch_size,
mode="train",
output_name=["buffer_A", "index_A"]),
BufferUpdate(input_name="fake_B",
buffer_size=50,
batch_size=batch_size,
mode="train",
output_name=["buffer_B", "index_B"]),
ModelSaver(model=g_AtoB, save_dir=save_dir, frequency=5),
ModelSaver(model=g_BtoA, save_dir=save_dir, frequency=5),
LRScheduler(model=g_AtoB, lr_fn=lr_schedule),
LRScheduler(model=g_BtoA, lr_fn=lr_schedule),
LRScheduler(model=d_A, lr_fn=lr_schedule),
LRScheduler(model=d_B, lr_fn=lr_schedule)
]
estimator = fe.Estimator(network=network,
pipeline=pipeline,
epochs=epochs,
traces=traces,
train_steps_per_epoch=train_steps_per_epoch)
return estimator
def get_estimator(data_dir=None,
epochs=12,
batch_size_per_gpu=4,
im_size=1344,
model_dir=tempfile.mkdtemp(),
train_steps_per_epoch=None,
eval_steps_per_epoch=None):
assert im_size % 32 == 0, "im_size must be a multiple of 32"
num_device = get_num_devices()
train_ds, val_ds = mscoco.load_data(root_dir=data_dir, load_masks=True)
batch_size = num_device * batch_size_per_gpu
pipeline = fe.Pipeline(
train_data=train_ds,
eval_data=val_ds,
test_data=val_ds,
ops=[
ReadImage(inputs="image", outputs="image"),
MergeMask(inputs="mask", outputs="mask"),
GetImageSize(inputs="image", outputs="imsize", mode="test"),
LongestMaxSize(max_size=im_size,
image_in="image",
mask_in="mask",
bbox_in="bbox",
bbox_params="coco"),
RemoveIf(fn=lambda x: len(x) == 0, inputs="bbox"),
PadIfNeeded(min_height=im_size,
min_width=im_size,
image_in="image",
mask_in="mask",
bbox_in="bbox",
bbox_params="coco",
border_mode=cv2.BORDER_CONSTANT,
value=0),
Sometimes(
HorizontalFlip(image_in="image",
mask_in="mask",
bbox_in="bbox",
bbox_params="coco",
mode="train")),
Resize(height=im_size // 4, width=im_size // 4,
image_in='mask'), # downscale mask for memory efficiency
Gt2Target(inputs=("mask", "bbox"),
outputs=("gt_match", "mask", "classes")),
Delete(keys="bbox"),
Delete(keys="image_id", mode="!test"),
Batch(batch_size=batch_size, pad_value=0)
],
num_process=8 * num_device)
init_lr = 1e-2 / 16 * batch_size
model = fe.build(
model_fn=SoloV2,
optimizer_fn=lambda x: torch.optim.SGD(x, lr=init_lr, momentum=0.9))
network = fe.Network(ops=[
Normalize(inputs="image",
outputs="image",
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225)),
Permute(inputs="image", outputs='image'),
ModelOp(model=model,
inputs="image",
outputs=("feat_seg", "feat_cls_list", "feat_kernel_list")),
LambdaOp(fn=lambda x: x,
inputs="feat_cls_list",
outputs=("cls1", "cls2", "cls3", "cls4", "cls5")),
LambdaOp(fn=lambda x: x,
inputs="feat_kernel_list",
outputs=("k1", "k2", "k3", "k4", "k5")),
Solov2Loss(0,
40,
inputs=("mask", "classes", "gt_match", "feat_seg", "cls1",
"k1"),
outputs=("l_c1", "l_s1")),
Solov2Loss(1,
36,
inputs=("mask", "classes", "gt_match", "feat_seg", "cls2",
"k2"),
outputs=("l_c2", "l_s2")),
Solov2Loss(2,
24,
inputs=("mask", "classes", "gt_match", "feat_seg", "cls3",
"k3"),
outputs=("l_c3", "l_s3")),
Solov2Loss(3,
16,
inputs=("mask", "classes", "gt_match", "feat_seg", "cls4",
"k4"),
outputs=("l_c4", "l_s4")),
Solov2Loss(4,
12,
inputs=("mask", "classes", "gt_match", "feat_seg", "cls5",
"k5"),
outputs=("l_c5", "l_s5")),
CombineLoss(inputs=("l_c1", "l_s1", "l_c2", "l_s2", "l_c3", "l_s3",
"l_c4", "l_s4", "l_c5", "l_s5"),
outputs=("total_loss", "cls_loss", "seg_loss")),
L2Regularizaton(inputs="total_loss",
outputs="total_loss_l2",
model=model,
beta=1e-5,
mode="train"),
UpdateOp(model=model, loss_name="total_loss_l2"),
PointsNMS(inputs="feat_cls_list", outputs="feat_cls_list",
mode="test"),
Predict(inputs=("feat_seg", "feat_cls_list", "feat_kernel_list"),
outputs=("seg_preds", "cate_scores", "cate_labels"),
mode="test")
])
train_steps_epoch = int(np.ceil(len(train_ds) / batch_size))
lr_schedule = {
1:
LRScheduler(
model=model,
lr_fn=lambda step: lr_schedule_warmup(step, init_lr=init_lr)),
2:
LRScheduler(
model=model,
lr_fn=lambda step: cosine_decay(step,
cycle_length=train_steps_epoch *
(epochs - 1),
init_lr=init_lr,
min_lr=init_lr / 100,
start=train_steps_epoch))
}
traces = [
EpochScheduler(lr_schedule),
COCOMaskmAP(data_dir=val_ds.root_dir,
inputs=("seg_preds", "cate_scores", "cate_labels",
"image_id", "imsize"),
mode="test"),
BestModelSaver(model=model, save_dir=model_dir, metric="total_loss")
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
monitor_names=("cls_loss", "seg_loss"),
train_steps_per_epoch=train_steps_per_epoch,
eval_steps_per_epoch=eval_steps_per_epoch)
return estimator