def get_estimator(epochs=50, batch_size=256, max_train_steps_per_epoch=None, save_dir=tempfile.mkdtemp()):
train_data, _ = mnist.load_data()
pipeline = fe.Pipeline(
train_data=train_data,
batch_size=batch_size,
ops=[
ExpandDims(inputs="x", outputs="x"),
Normalize(inputs="x", outputs="x", mean=1.0, std=1.0, max_pixel_value=127.5),
LambdaOp(fn=lambda: np.random.normal(size=[100]).astype('float32'), outputs="z")
])
gen_model = fe.build(model_fn=generator, optimizer_fn=lambda: tf.optimizers.Adam(1e-4))
disc_model = fe.build(model_fn=discriminator, optimizer_fn=lambda: tf.optimizers.Adam(1e-4))
network = fe.Network(ops=[
ModelOp(model=gen_model, inputs="z", outputs="x_fake"),
ModelOp(model=disc_model, inputs="x_fake", outputs="fake_score"),
GLoss(inputs="fake_score", outputs="gloss"),
UpdateOp(model=gen_model, loss_name="gloss"),
ModelOp(inputs="x", model=disc_model, outputs="true_score"),
DLoss(inputs=("true_score", "fake_score"), outputs="dloss"),
UpdateOp(model=disc_model, loss_name="dloss")
])
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=ModelSaver(model=gen_model, save_dir=save_dir, frequency=5),
max_train_steps_per_epoch=max_train_steps_per_epoch)
return estimator
def test_repeat_fn_exterior_value(self):
add_op = LambdaOp(inputs='x', outputs=('x', 'y'), fn=lambda x: (x + 1, x * x), mode='eval')
repeat_op = Repeat(add_op, repeat=lambda y, z: y + z < 25)
with self.subTest('Check op inputs'):
self.assertListEqual(repeat_op.inputs, ['x', 'z'])
with self.subTest('Check op outputs'):
self.assertListEqual(repeat_op.outputs, ['x', 'y'])
with self.subTest('Check op mode'):
self.assertSetEqual(repeat_op.mode, {'eval'})
output = repeat_op.forward(data=[np.ones([1]), 10 + np.ones([1])], state={"mode": "eval"})
with self.subTest('Check output type'):
self.assertEqual(type(output), list)
with self.subTest('Check output value (x)'):
self.assertEqual(5, output[0])
with self.subTest('Check output value (y)'):
self.assertEqual(16, output[1])
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: max(512 // size, 4) * get_num_devices() if size <= 512 else 2 * 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),
ChannelTranspose(inputs=["x", "x_low_res"], outputs=["x", "x_low_res"]),
LambdaOp(fn=lambda: np.random.normal(size=[512]).astype('float32'), outputs="z")
])
fade_in_alpha = torch.tensor(1.0)
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 x: Adam(x, lr=0.001, betas=(0.0, 0.99), eps=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 x: Adam(x, lr=0.001, betas=(0.0, 0.99), eps=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(epochs=20, batch_size=128, max_train_steps_per_epoch=None, save_dir=tempfile.mkdtemp()):
# Dataset Creation
(x_train, y_train), (x_eval, y_eval) = tf.keras.datasets.mnist.load_data()
x_eval0, y_eval0 = x_eval[np.where((y_eval == 1))], np.ones(y_eval[np.where((y_eval == 1))].shape)
x_eval1, y_eval1 = x_eval[np.where((y_eval != 1))], y_eval[np.where((y_eval != 1))]
# Ensuring outliers comprise 50% of the dataset
index = np.random.choice(x_eval1.shape[0], int(x_eval0.shape[0]), replace=False)
x_eval1, y_eval1 = x_eval1[index], np.zeros(y_eval1[index].shape)
x_train, y_train = x_train[np.where((y_train == 1))], np.zeros(y_train[np.where((y_train == 1))].shape)
train_data = fe.dataset.NumpyDataset({"x": x_train, "y": y_train})
x_eval, y_eval = np.concatenate([x_eval0, x_eval1]), np.concatenate([y_eval0, y_eval1])
eval_data = fe.dataset.NumpyDataset({"x": x_eval, "y": y_eval})
pipeline = fe.Pipeline(
train_data=train_data,
eval_data=eval_data,
batch_size=batch_size,
ops=[
ExpandDims(inputs="x", outputs="x"),
Normalize(inputs="x", outputs="x", mean=1.0, std=1.0, max_pixel_value=127.5),
LambdaOp(fn=lambda x: x + np.random.normal(loc=0.0, scale=0.155, size=(28, 28, 1)),
inputs="x",
outputs="x_w_noise",
mode="train")
])
recon_model = fe.build(model_fn=reconstructor,
optimizer_fn=lambda: tf.optimizers.RMSprop(2e-4),
model_name="reconstructor")
disc_model = fe.build(model_fn=discriminator,
optimizer_fn=lambda: tf.optimizers.RMSprop(1e-4),
model_name="discriminator")
network = fe.Network(ops=[
ModelOp(model=recon_model, inputs="x_w_noise", outputs="x_fake", mode="train"),
ModelOp(model=recon_model, inputs="x", outputs="x_fake", mode="eval"),
ModelOp(model=disc_model, inputs="x_fake", outputs="fake_score"),
ModelOp(model=disc_model, inputs="x", outputs="true_score"),
RLoss(inputs=("fake_score", "x_fake", "x"), outputs="rloss"),
UpdateOp(model=recon_model, loss_name="rloss"),
DLoss(inputs=("true_score", "fake_score"), outputs="dloss"),
UpdateOp(model=disc_model, loss_name="dloss"),
])
traces = [
F1AUCScores(true_key="y", pred_key="fake_score", mode="eval", output_name=["auc_score", "f1_score"]),
BestModelSaver(model=recon_model, save_dir=save_dir, metric='f1_score', save_best_mode='max'),
BestModelSaver(model=disc_model, save_dir=save_dir, metric='f1_score', save_best_mode='max'),
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch,
log_steps=50)
return estimator
def test_batch_forward(self):
op = LambdaOp(fn=np.sum)
data = tf.convert_to_tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
result = op.forward_batch(data=[data], state={})
ans = np.array([6, 15, 24], dtype=np.float32)
self.assertTrue(np.array_equal(result, ans))
def test_multi_input(self):
op = LambdaOp(fn=np.reshape)
data = op.forward(data=[np.array([1, 2, 3, 4]), (2, 2)], state={})
self.assertTrue(is_equal(data, np.array([[1, 2], [3, 4]])))
def test_single_input(self):
op = LambdaOp(fn=np.sum)
data = op.forward(data=[[1, 2, 3]], state={})
self.assertEqual(data, 6)
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:0" 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 test_batch_forward(self):
op = LambdaOp(fn=np.sum)
data = tf.convert_to_tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
result = op.forward_batch(data=[data], state={})
self.assertEqual(result, 45)
