def test_pytorch_weight_decay_vs_l2(self):
# Get Data
train_data, _ = mnist.load_data()
t_d = train_data.split(128)
# Initializing models
pytorch_wd = fe.build(model_fn=MyNet_torch,
optimizer_fn=lambda x: torch.optim.SGD(params=x, lr=0.01, weight_decay=self.beta))
pytorch_l2 = fe.build(model_fn=MyNet_torch, optimizer_fn=lambda x: torch.optim.SGD(params=x, lr=0.01))
# Initialize pipeline
pipeline = fe.Pipeline(train_data=t_d,
batch_size=128,
ops=[ExpandDims(inputs="x", outputs="x", axis=0), Minmax(inputs="x", outputs="x")])
# Define the two pytorch networks
network_weight_decay = fe.Network(ops=[
ModelOp(model=pytorch_wd, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
UpdateOp(model=pytorch_wd, loss_name="ce")
])
network_l2 = fe.Network(ops=[
ModelOp(model=pytorch_l2, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
L2Regularizaton(inputs="ce", outputs="l2", model=pytorch_l2, beta=self.beta),
UpdateOp(model=pytorch_l2, loss_name="l2")
])
# defining traces
traces = [Accuracy(true_key="y", pred_key="y_pred")]
# Setting up estimators
estimator_wd = fe.Estimator(pipeline=pipeline,
network=network_weight_decay,
epochs=1,
traces=traces,
train_steps_per_epoch=1)
estimator_l2 = fe.Estimator(pipeline=pipeline,
network=network_l2,
epochs=1,
traces=traces,
train_steps_per_epoch=1)
# Training
print('********************************Pytorch weight decay training************************************')
estimator_wd.fit()
print()
print('********************************Pytorch L2 Regularization training************************************')
estimator_l2.fit()
# testing weights
count = 0
for wt, l2 in zip(pytorch_wd.parameters(), pytorch_l2.parameters()):
if ((wt - l2).abs()).sum() < torch.tensor(10**-6):
count += 1
self.assertTrue(count == 6)
def get_estimator(epochs=24, batch_size=128, lr_epochs=100, max_train_steps_per_epoch=None,
save_dir=tempfile.mkdtemp()):
# step 1: prepare dataset
train_data, test_data = load_data()
pipeline = fe.Pipeline(
train_data=train_data,
eval_data=test_data,
batch_size=batch_size,
ops=[
Normalize(inputs="x", outputs="x", mean=(0.4914, 0.4822, 0.4465), std=(0.2471, 0.2435, 0.2616)),
PadIfNeeded(min_height=40, min_width=40, image_in="x", image_out="x", mode="train"),
RandomCrop(32, 32, image_in="x", image_out="x", mode="train"),
Sometimes(HorizontalFlip(image_in="x", image_out="x", mode="train")),
CoarseDropout(inputs="x", outputs="x", mode="train", max_holes=1),
ChannelTranspose(inputs="x", outputs="x"),
Onehot(inputs="y", outputs="y", mode="train", num_classes=10, label_smoothing=0.2)
])
# step 2: prepare network
model = fe.build(model_fn=ResNet9, optimizer_fn="sgd")
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
UpdateOp(model=model, loss_name="ce")
])
# get the max learning rate
lr_max = search_max_lr(pipeline=pipeline, model=model, network=network, epochs=lr_epochs)
lr_min = lr_max / 40
print(f"The maximum LR: {lr_max}, and minimun LR: {lr_min}")
mid_step = int(epochs * 0.45 * len(train_data) / batch_size)
end_step = int(epochs * len(train_data) / batch_size)
# reinitialize the model
model = fe.build(model_fn=ResNet9, optimizer_fn="sgd")
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
UpdateOp(model=model, loss_name="ce")
])
# step 3: prepare estimator
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
BestModelSaver(model=model, save_dir=save_dir, metric="accuracy", save_best_mode="max"),
LRScheduler(model=model, lr_fn=lambda step: super_schedule(step, lr_max, lr_min, mid_step, end_step))
]
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=2,
batch_size=32,
save_dir=tempfile.mkdtemp()):
# step 1
train_data, eval_data = mnist.load_data()
test_data = eval_data.split(0.5)
pipeline = fe.Pipeline(train_data=train_data,
eval_data=eval_data,
test_data=test_data,
batch_size=batch_size,
ops=[ExpandDims(inputs="x", outputs="x"), Minmax(inputs="x", outputs="x")])
# step 2
model = fe.build(model_fn=LeNet, optimizer_fn="adam")
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
UpdateOp(model=model, loss_name="ce")
])
# step 3
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
BestModelSaver(model=model, save_dir=save_dir, metric="accuracy", save_best_mode="max"),
LRScheduler(model=model, lr_fn=lambda step: cosine_decay(step, cycle_length=3750, init_lr=1e-3))
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces)
return estimator
def get_estimator(epochs=10, batch_size=32, extend_ds=False):
train_data, eval_data = cifair10.load_data()
if extend_ds:
train_data = ExtendDataset(dataset=train_data,
spoof_length=len(train_data) * 2)
epochs //= 2
pipeline = fe.Pipeline(train_data=train_data,
eval_data=eval_data,
batch_size=batch_size,
ops=[Normalize(inputs="x", outputs="x")])
# step 2
model = fe.build(model_fn=lambda: LeNet(input_shape=(32, 32, 3)),
optimizer_fn="adam")
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
UpdateOp(model=model, loss_name="ce")
])
# step 3
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces)
return estimator
def get_estimator(epochs=10, batch_size=32, alpha=1.0, warmup=0, model_dir=tempfile.mkdtemp()):
(x_train, y_train), (x_eval, y_eval) = tf.keras.datasets.cifar10.load_data()
data = {"train": {"x": x_train, "y": y_train}, "eval": {"x": x_eval, "y": y_eval}}
num_classes = 10
pipeline = fe.Pipeline(batch_size=batch_size, data=data, ops=Minmax(inputs="x", outputs="x"))
model = FEModel(model_def=lambda: LeNet(input_shape=x_train.shape[1:], classes=num_classes),
model_name="LeNet",
optimizer="adam")
mixup_map = {warmup: MixUpBatch(inputs="x", outputs=["x", "lambda"], alpha=alpha, mode="train")}
mixup_loss = {
0: SparseCategoricalCrossentropy(y_true="y", y_pred="y_pred", mode="train"),
warmup: MixUpLoss(KerasCrossentropy(), lam="lambda", y_true="y", y_pred="y_pred", mode="train")
}
network = fe.Network(ops=[
Scheduler(mixup_map),
ModelOp(inputs="x", model=model, outputs="y_pred"),
Scheduler(mixup_loss),
SparseCategoricalCrossentropy(y_true="y", y_pred="y_pred", mode="eval")
])
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
ConfusionMatrix(true_key="y", pred_key="y_pred", num_classes=num_classes),
ModelSaver(model_name="LeNet", save_dir=model_dir, save_best=True)
]
estimator = fe.Estimator(network=network, pipeline=pipeline, epochs=epochs, traces=traces)
return estimator
def get_estimator(epochs=2, batch_size=32):
# step 1
train_data, eval_data = mnist.load_data()
pipeline = fe.Pipeline(train_data=train_data,
eval_data=eval_data,
batch_size=batch_size,
ops=[
ExpandDims(inputs="x", outputs="x"),
Minmax(inputs="x", outputs="x")
])
# step 2
model = fe.build(model_fn=LeNet, optimizer_fn="adam")
network = fe.Network(ops=[
ModelOp(model=model,
inputs="x",
outputs=["y_pred", "feature_vector"],
intermediate_layers='dense'),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
CustomLoss(inputs=("feature_vector", "feature_selected"),
outputs="feature_loss"),
LambdaOp(fn=lambda x, y: x + y,
inputs=("ce", "feature_loss"),
outputs="total_loss"),
UpdateOp(model=model, loss_name="total_loss")
])
# step 3
traces = [
MemoryBank(inputs=("feature_vector", "y"), outputs="feature_selected")
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces)
return estimator
def get_estimator(epochs=10, batch_size=32, train_steps_per_epoch=None):
train_data, eval_data = cifair10.load_data()
pipeline = fe.Pipeline(train_data=train_data,
eval_data=eval_data,
batch_size=batch_size,
ops=[Normalize(inputs="x", outputs="x")])
# step 2
model = fe.build(model_fn=lambda: LeNet(input_shape=(32, 32, 3)),
optimizer_fn="adam")
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
UpdateOp(model=model, loss_name="ce")
])
# step 3
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
train_steps_per_epoch=train_steps_per_epoch)
return estimator
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 get_estimator(style_img_path=None,
data_path=None,
style_weight=5.0,
content_weight=1.0,
tv_weight=1e-4,
steps_per_epoch=None,
validation_steps=None,
model_dir=tempfile.mkdtemp()):
train_csv, _, path = load_data(data_path, load_object=False)
if style_img_path is None:
style_img_path = tf.keras.utils.get_file(
'kandinsky.jpg',
'https://storage.googleapis.com/download.tensorflow.org/example_images/Vassily_Kandinsky%2C_1913_'
'-_Composition_7.jpg')
style_img = cv2.imread(style_img_path)
assert (style_img is not None), "Invalid style reference image"
tfr_save_dir = os.path.join(path, 'tfrecords')
style_img = (style_img.astype(np.float32) / 127.5) / 127.5
style_img_t = tf.convert_to_tensor(np.expand_dims(style_img, axis=0))
writer = RecordWriter(train_data=train_csv,
save_dir=tfr_save_dir,
ops=[
ImageReader(inputs="image",
parent_path=path,
outputs="image"),
Resize(inputs="image",
target_size=(256, 256),
outputs="image")
])
pipeline = fe.Pipeline(batch_size=4,
data=writer,
ops=[Rescale(inputs="image", outputs="image")])
model = fe.build(model_def=styleTransferNet,
model_name="style_transfer_net",
loss_name="loss",
optimizer=tf.keras.optimizers.Adam(1e-3))
network = fe.Network(ops=[
ModelOp(inputs="image", model=model, outputs="image_out"),
ExtractVGGFeatures(inputs=lambda: style_img_t, outputs="y_style"),
ExtractVGGFeatures(inputs="image", outputs="y_content"),
ExtractVGGFeatures(inputs="image_out", outputs="y_pred"),
StyleContentLoss(style_weight=style_weight,
content_weight=content_weight,
tv_weight=tv_weight,
inputs=('y_pred', 'y_style', 'y_content',
'image_out'),
outputs='loss')
])
estimator = fe.Estimator(network=network,
pipeline=pipeline,
epochs=2,
steps_per_epoch=steps_per_epoch,
validation_steps=validation_steps,
traces=ModelSaver(model_name="style_transfer_net",
save_dir=model_dir))
return estimator
def get_estimator(batch_size=32, epochs=25, model_dir=tempfile.mkdtemp()):
# load CUB200 dataset.
csv_path, path = cub200.load_data()
writer = RecordWriter(
save_dir=os.path.join(path, "FEdata"),
train_data=csv_path,
validation_data=0.2,
ops=[
ImageReader(inputs='image', parent_path=path),
Resize(target_size=(128, 128), keep_ratio=True, outputs='image'),
MatReader(inputs='annotation', parent_path=path),
SelectDictKey(),
Resize((128, 128), keep_ratio=True),
Reshape(shape=(128, 128, 1), outputs="annotation")
])
# data pipeline
pipeline = fe.Pipeline(batch_size=batch_size, data=writer, ops=Minmax(inputs='image', outputs='image'))
# Network
model = FEModel(model_def=UNet, model_name="unet_cub", optimizer=tf.optimizers.Adam())
network = fe.Network(ops=[
ModelOp(inputs='image', model=model, outputs='mask_pred'),
BinaryCrossentropy(y_true='annotation', y_pred='mask_pred')
])
# estimator
traces = [
Dice(true_key="annotation", pred_key='mask_pred'),
ModelSaver(model_name="unet_cub", save_dir=model_dir, save_best=True)
]
estimator = fe.Estimator(network=network, pipeline=pipeline, traces=traces, epochs=epochs, log_steps=50)
return estimator
def get_estimator(batch_size=256, epochs=50, model_dir=tempfile.mkdtemp()):
# prepare data
(x_train, _), (_, _) = tf.keras.datasets.mnist.load_data()
data = {"train": {"x": np.expand_dims(x_train, -1)}}
pipeline = fe.Pipeline(batch_size=batch_size,
data=data,
ops=Myrescale(inputs="x", outputs="x"))
# prepare model
g_femodel = FEModel(model_def=make_generator_model,
model_name="gen",
loss_name="gloss",
optimizer=tf.optimizers.Adam(1e-4))
d_femodel = FEModel(model_def=make_discriminator_model,
model_name="disc",
loss_name="dloss",
optimizer=tf.optimizers.Adam(1e-4))
network = fe.Network(ops=[
ModelOp(inputs=lambda: tf.random.normal([batch_size, 100]),
model=g_femodel),
ModelOp(model=d_femodel, outputs="pred_fake"),
ModelOp(inputs="x", model=d_femodel, outputs="pred_true"),
GLoss(inputs=("pred_fake"), outputs="gloss"),
DLoss(inputs=("pred_true", "pred_fake"), outputs="dloss")
])
# prepare estimator
traces = [ModelSaver(model_name='gen', save_dir=model_dir, save_freq=5)]
estimator = fe.Estimator(network=network,
pipeline=pipeline,
epochs=epochs,
traces=traces)
return estimator
def build_network(self):
"""
Define the FastEstimator network flow.
Args:
None
Returns:
network: KerasNetwork object
"""
epsilon = 0.04
network = fe.Network(ops=[
Watch(inputs="x"),
ModelOp(model=self.model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="base_ce"),
FGSM(
data="x", loss="base_ce", outputs="x_adverse",
epsilon=epsilon),
ModelOp(model=self.model, inputs="x_adverse",
outputs="y_pred_adv"),
CrossEntropy(inputs=("y_pred_adv", "y"), outputs="adv_ce"),
Average(inputs=("base_ce", "adv_ce"), outputs="avg_ce"),
UpdateOp(model=self.model, loss_name="avg_ce")
])
return network
def test_network_transform_one_layer_model_torch(self):
model = fe.build(model_fn=OneLayerTorchModel, optimizer_fn="adam")
weight = get_torch_one_layer_model_weight(model)
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
MeanSquaredError(inputs=("y_pred", "y"), outputs="ce"),
UpdateOp(model=model, loss_name="ce")
],
pops=PlusOneNumpyOp(inputs="y_pred",
outputs="y_pred_processed"))
batch = {
"x": np.array([[
1,
1,
1,
]], dtype=np.float32),
"y": np.array([[1]], dtype=np.float32)
}
batch = network.transform(data=batch, mode="train")
with self.subTest("output y_pred check"):
ans = np.array([[6]], dtype=np.float32) # 1*1 + 1*2 + 1*3
self.assertTrue(np.array_equal(batch["y_pred"].numpy(), ans))
with self.subTest("postprocessing y_pred check"):
ans = np.array([[7]], dtype=np.float32) # 1*1 + 1*2 + 1*3 + 1
self.assertTrue(np.array_equal(batch["y_pred_processed"], ans))
with self.subTest("output ce check"):
self.assertEqual(batch["ce"].numpy(), 25) # (6-1)^2
with self.subTest("check whether model weight changed"):
weight2 = get_torch_one_layer_model_weight(model)
self.assertFalse(is_equal(weight, weight2))
def test_network_transform_lenet_torch(self):
model = fe.build(
model_fn=LeNetTorch,
optimizer_fn=lambda x: torch.optim.Adam(params=x, lr=1.0))
weight = get_torch_lenet_model_weight(model)
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
UpdateOp(model=model, loss_name="ce")
])
batch = {
"x": np.ones((1, 1, 28, 28), dtype=np.float32),
"y": np.array([[1]], dtype=np.float32)
}
batch = network.transform(data=batch, mode="train")
with self.subTest("output y_pred check"):
self.assertTrue("y_pred" in batch.keys())
self.assertIsNotNone(batch["y_pred"])
with self.subTest("output ce check"):
self.assertTrue("ce" in batch.keys())
self.assertIsNotNone(batch["ce"])
with self.subTest("check whether model weight changed"):
weight2 = get_torch_lenet_model_weight(model)
self.assertFalse(is_equal(weight, weight2))
def finetune_model(model, epochs, batch_size, max_train_steps_per_epoch, save_dir):
train_data, test_data = load_data()
train_data = train_data.split(0.1)
pipeline = fe.Pipeline(train_data=train_data,
eval_data=test_data,
batch_size=batch_size,
ops=[
ToFloat(inputs="x", outputs="x"),
])
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=["y_pred", "y"], outputs="ce"),
UpdateOp(model=model, loss_name="ce")
])
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch)
estimator.fit()
def run_test(mixed_precision, merge_grad, gradient):
lr = 0.1
pipeline = fe.Pipeline(train_data=self.train_data,
batch_size=4,
ops=[ExpandDims(inputs="x", outputs="x"), Minmax(inputs="x", outputs="x")])
model = fe.build(model_fn=LeNet_tf,
optimizer_fn=lambda: tf.optimizers.SGD(lr),
mixed_precision=mixed_precision)
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
GradientOp(model=model, finals="ce", outputs="grad"),
UpdateOp(model=model, loss_name="ce", gradients=gradient, merge_grad=merge_grad),
])
traces = [
CheckNetworkWeight(model=model,
grad_key="grad",
merge_grad=merge_grad,
test_self=self,
lrs=lr,
framework="tf")
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=2,
traces=traces,
train_steps_per_epoch=2)
estimator.fit(warmup=False)
def get_estimator(epochs=2, batch_size=32, model_dir=tempfile.mkdtemp()):
# step 1. prepare data
(x_train, y_train), (x_eval, y_eval) = tf.keras.datasets.mnist.load_data()
train_data = {"x": np.expand_dims(x_train, -1), "y": y_train}
eval_data = {"x": np.expand_dims(x_eval, -1), "y": y_eval}
data = {"train": train_data, "eval": eval_data}
pipeline = fe.Pipeline(batch_size=batch_size,
data=data,
ops=Minmax(inputs="x", outputs="x"))
# step 2. prepare model
model = fe.FEModel(model_def=LeNet, model_name="lenet", optimizer="adam")
network = fe.Network(ops=[
ModelOp(inputs="x", model=model, outputs="y_pred"),
SparseCategoricalCrossentropy(inputs=("y", "y_pred"))
])
# step 3.prepare estimator
traces = [
Accuracy(true_key="y", pred_key="y_pred", output_name='acc'),
ModelSaver(model_name="lenet", save_dir=model_dir, save_best=True)
]
estimator = fe.Estimator(network=network,
pipeline=pipeline,
epochs=epochs,
traces=traces)
return estimator
def get_estimator(epochs=10, batch_size=50, epsilon=0.04, save_dir=tempfile.mkdtemp()):
train_data, eval_data = cifar10.load_data()
test_data = eval_data.split(0.5)
pipeline = fe.Pipeline(
train_data=train_data,
eval_data=eval_data,
test_data=test_data,
batch_size=batch_size,
ops=[Normalize(inputs="x", outputs="x", mean=(0.4914, 0.4822, 0.4465), std=(0.2471, 0.2435, 0.2616))])
clean_model = fe.build(model_fn=lambda: LeNet(input_shape=(32, 32, 3)),
optimizer_fn="adam",
model_name="clean_model")
clean_network = fe.Network(ops=[
Watch(inputs="x"),
ModelOp(model=clean_model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="base_ce"),
FGSM(data="x", loss="base_ce", outputs="x_adverse", epsilon=epsilon, mode="!train"),
ModelOp(model=clean_model, inputs="x_adverse", outputs="y_pred_adv", mode="!train"),
UpdateOp(model=clean_model, loss_name="base_ce")
])
clean_traces = [
Accuracy(true_key="y", pred_key="y_pred", output_name="clean_accuracy"),
Accuracy(true_key="y", pred_key="y_pred_adv", output_name="adversarial_accuracy"),
BestModelSaver(model=clean_model, save_dir=save_dir, metric="base_ce", save_best_mode="min"),
]
clean_estimator = fe.Estimator(pipeline=pipeline,
network=clean_network,
epochs=epochs,
traces=clean_traces,
log_steps=1000)
return clean_estimator
def _build_estimator(model: Union[tf.keras.Model, torch.nn.Module], trace: Traceability, axis: int = -1):
train_data, eval_data = mnist.load_data()
test_data = eval_data.split(0.5)
batch_size = 32
pipeline = fe.Pipeline(train_data=train_data,
eval_data=eval_data,
test_data=test_data,
batch_size=batch_size,
ops=[ExpandDims(inputs="x", outputs="x", axis=axis), Minmax(inputs="x", outputs="x")])
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
UpdateOp(model=model, loss_name="ce")
])
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
LRScheduler(model=model, lr_fn=lambda step: cosine_decay(step, cycle_length=3750, init_lr=1e-3)),
trace
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=1,
traces=traces,
max_train_steps_per_epoch=1,
max_eval_steps_per_epoch=None)
fake_data = tf.ones(shape=(batch_size, 28, 28, 1)) if axis == -1 else torch.ones(size=(batch_size, 1, 28, 28))
model.fe_input_spec = FeInputSpec(fake_data, model)
return estimator
def test_mode_ds_id_interaction(self):
train_data, eval_data = mnist.load_data()
test_data = eval_data.split(0.5)
pipeline = fe.Pipeline(train_data=train_data,
eval_data=eval_data,
test_data=test_data,
batch_size=32,
ops=[
ExpandDims(inputs="x", outputs="x"),
Minmax(inputs="x", outputs="x")
])
model = fe.build(model_fn=LeNet, optimizer_fn="adam")
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce", ds_id="ds_1")
])
pipeline_data = pipeline.transform(data=train_data[0], mode="train")
data1 = network.transform(data=pipeline_data,
mode="infer",
ds_id="ds_1")
assert "ce" not in data1
data2 = network.transform(data=pipeline_data,
mode="infer",
ds_id="ds_2")
assert "ce" not in data2
def get_estimator(batch_size=4, epochs=25, model_dir=tempfile.mkdtemp()):
csv_path, path = montgomery.load_data()
writer = RecordWriter(save_dir=os.path.join(path, "FEdata"),
train_data=csv_path,
validation_data=0.2,
ops=[
ImageReader(grey_scale=True,
inputs="image",
parent_path=path,
outputs="image"),
ImageReader(grey_scale=True,
inputs="mask_left",
parent_path=path,
outputs="mask_left"),
ImageReader(grey_scale=True,
inputs="mask_right",
parent_path=path,
outputs="mask_right"),
CombineLeftRightMask(inputs=("mask_left",
"mask_right")),
Resize(target_size=(512, 512)),
Reshape(shape=(512, 512, 1), outputs="mask"),
Resize(inputs="image", target_size=(512, 512)),
Reshape(shape=(512, 512, 1), outputs="image"),
],
write_feature=["image", "mask"])
pipeline = fe.Pipeline(batch_size=batch_size,
data=writer,
ops=[
Augmentation2D(inputs=["image", "mask"],
outputs=["image", "mask"],
mode="train",
rotation_range=10,
flip_left_right=True),
Minmax(inputs="image", outputs="image"),
Minmax(inputs="mask", outputs="mask")
])
model = FEModel(model_def=lambda: UNet(input_size=(512, 512, 1)),
model_name="lungsegmentation",
optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001))
network = fe.Network(ops=[
ModelOp(inputs="image", model=model, outputs="pred_segment"),
BinaryCrossentropy(y_true="mask", y_pred="pred_segment")
])
traces = [
Dice(true_key="mask", pred_key="pred_segment"),
ModelSaver(model_name="lungsegmentation",
save_dir=model_dir,
save_best=True)
]
estimator = fe.Estimator(network=network,
pipeline=pipeline,
epochs=epochs,
log_steps=20,
traces=traces)
return estimator
def run_test(mixed_precision, merge_grad, gradient):
lr = 0.1
lr2 = 0.01
pipeline = fe.Pipeline(train_data=self.train_data,
batch_size=4,
ops=[ExpandDims(inputs="x", outputs="x", axis=0), Minmax(inputs="x", outputs="x")])
optimizer_fn = RepeatScheduler(
[lambda x: torch.optim.SGD(params=x, lr=lr), lambda x: torch.optim.SGD(params=x, lr=lr2)])
model = fe.build(model_fn=LeNet_torch, optimizer_fn=optimizer_fn, mixed_precision=mixed_precision)
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
GradientOp(model=model, finals="ce", outputs="grad"),
UpdateOp(model=model, loss_name="ce", gradients=gradient, merge_grad=merge_grad),
])
traces = [
CheckNetworkWeight(model=model,
grad_key="grad",
merge_grad=merge_grad,
test_self=self,
framework="torch",
lrs=[lr, lr2, lr, lr2],
work_intervals=[[1, 2], [2, 3], [3, 4], [4, 5]])
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=4,
traces=traces,
train_steps_per_epoch=2)
estimator.fit(warmup=False)
def test_estimator_configure_loader_torch_data_loader_tf_model(self):
loader = get_sample_torch_dataloader()
pipeline = fe.Pipeline(train_data=loader)
model = fe.build(model_fn=LeNetTf, optimizer_fn="adam")
network = fe.Network(ops=[
ModelOp(model=model, inputs="x_out", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
UpdateOp(model=model, loss_name="ce")
])
est = fe.Estimator(pipeline=pipeline,
network=network,
max_train_steps_per_epoch=3,
epochs=1)
est.system.mode = "train"
new_loader = est._configure_loader(loader)
with self.subTest("check loader type"):
strategy = tf.distribute.get_strategy()
if isinstance(strategy, tf.distribute.MirroredStrategy):
self.assertIsInstance(new_loader,
tf.distribute.DistributedDataset)
else:
self.assertIsInstance(new_loader, tf.data.Dataset)
with self.subTest("max_train_steps_per_epoch=3"):
iterator = iter(new_loader)
for i in range(3):
batch = next(iterator)
with self.assertRaises(StopIteration):
batch = next(iterator)
def get_estimator(max_words=10000,
max_len=500,
epochs=10,
batch_size=64,
train_steps_per_epoch=None,
eval_steps_per_epoch=None,
save_dir=tempfile.mkdtemp()):
# step 1. prepare data
train_data, eval_data = imdb_review.load_data(max_len, max_words)
pipeline = fe.Pipeline(train_data=train_data,
eval_data=eval_data,
batch_size=batch_size,
ops=Reshape(1, inputs="y", outputs="y"))
# step 2. prepare model
model = fe.build(model_fn=lambda: ReviewSentiment(max_words=max_words), optimizer_fn="adam")
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="loss"),
UpdateOp(model=model, loss_name="loss")
])
# step 3.prepare estimator
traces = [Accuracy(true_key="y", pred_key="y_pred"), BestModelSaver(model=model, save_dir=save_dir)]
estimator = fe.Estimator(network=network,
pipeline=pipeline,
epochs=epochs,
traces=traces,
train_steps_per_epoch=train_steps_per_epoch,
eval_steps_per_epoch=eval_steps_per_epoch)
return estimator
def instantiate_system():
system = sample_system_object()
model = fe.build(model_fn=fe.architecture.pytorch.LeNet, optimizer_fn='adam', model_name='tf')
system.network = fe.Network(ops=[
ModelOp(model=model, inputs="x_out", outputs="y_pred"),
SuperLoss(CrossEntropy(inputs=['y_pred', 'y'], outputs='ce'))
])
return system
def instantiate_system():
system = sample_system_object_torch()
model = fe.build(model_fn=fe.architecture.pytorch.LeNet, optimizer_fn='adam', model_name='torch')
system.network = fe.Network(ops=[
fe.op.tensorop.meta.Sometimes(TestTensorOp(inputs="x_out", outputs="x_out", mode="train", var=1)),
ModelOp(model=model, inputs="x_out", outputs="y_pred")
])
return system
def test_estimator_warmup_tf_dataset_torch_model_smoke(self):
loader = get_sample_tf_dataset(expand_axis=1)
pipeline = fe.Pipeline(train_data=loader) # "x", "y"
model = fe.build(model_fn=LeNetTorch, optimizer_fn="adam")
network = fe.Network(ops=[ModelOp(model=model, inputs="x", outputs="y_pred")])
est = fe.Estimator(pipeline=pipeline, network=network, epochs=1, traces=[Trace(inputs="y_pred")])
est._prepare_traces(run_modes={"train", "eval"})
est._warmup()
self.assertTrue(True)
def test_qms(self):
test_data = get_sample_tf_dataset()
pipeline = fe.Pipeline(test_data=test_data)
model = fe.build(model_fn=one_layer_tf_model, optimizer_fn="adam")
network = fe.Network(
ops=[ModelOp(model=model, inputs="x", outputs="y_pred")])
test_title = "Integration Test of QMSTest"
temp_dir = tempfile.mkdtemp()
json_output = os.path.join(temp_dir, "test.json")
doc_output = os.path.join(temp_dir, "summary.docx")
test_descriptions = [
"first result is greater than 0", "second result is greater than 0"
]
traces = [
PassTrace(inputs="y_pred", mode="test"),
QMSTest(
test_descriptions=test_descriptions,
test_criterias=[
lambda y_pred: y_pred[0][0][0] > 0, # 1*1 + 1*2 + 1*3 > 0
lambda y_pred: y_pred[0][1][0] >
0, # 1*1 + 2*2 + 3*(-3) > 0
],
test_title=test_title,
json_output=json_output,
doc_output=doc_output)
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=1,
traces=traces)
with patch("fastestimator.trace.io.qms.json.dump") as fake_dump, \
patch("fastestimator.trace.io.qms._QMSDocx") as fake_qms:
estimator.test()
(json_summary, json_fp), _ = fake_dump.call_args
with self.subTest("check json summary dict"):
ans = {
"title":
test_title,
"stories": [{
"description": test_descriptions[0],
"passed": "True"
}, {
"description": test_descriptions[1],
"passed": "False"
}]
}
self.assertTrue(is_equal(json_summary, ans))
with self.subTest("check json summary stored path"):
self.assertTrue(is_equal(json_fp.name, json_output))
with self.subTest("check call the _QMSDocx correctly"):
ans = (1, 1) # (total_pass, total_fail)
self.assertEqual(fake_qms.call_args[0], ans)
def pretrain_model(epochs, batch_size, max_train_steps_per_epoch, save_dir):
# step 1: prepare dataset
train_data, test_data = load_data()
pipeline = fe.Pipeline(
train_data=train_data,
batch_size=batch_size,
ops=[
PadIfNeeded(min_height=40, min_width=40, image_in="x", image_out="x"),
# augmentation 1
RandomCrop(32, 32, image_in="x", image_out="x_aug"),
Sometimes(HorizontalFlip(image_in="x_aug", image_out="x_aug"), prob=0.5),
Sometimes(
ColorJitter(inputs="x_aug", outputs="x_aug", brightness=0.8, contrast=0.8, saturation=0.8, hue=0.2),
prob=0.8),
Sometimes(ToGray(inputs="x_aug", outputs="x_aug"), prob=0.2),
Sometimes(GaussianBlur(inputs="x_aug", outputs="x_aug", blur_limit=(3, 3), sigma_limit=(0.1, 2.0)),
prob=0.5),
ToFloat(inputs="x_aug", outputs="x_aug"),
# augmentation 2
RandomCrop(32, 32, image_in="x", image_out="x_aug2"),
Sometimes(HorizontalFlip(image_in="x_aug2", image_out="x_aug2"), prob=0.5),
Sometimes(
ColorJitter(inputs="x_aug2", outputs="x_aug2", brightness=0.8, contrast=0.8, saturation=0.8, hue=0.2),
prob=0.8),
Sometimes(ToGray(inputs="x_aug2", outputs="x_aug2"), prob=0.2),
Sometimes(GaussianBlur(inputs="x_aug2", outputs="x_aug2", blur_limit=(3, 3), sigma_limit=(0.1, 2.0)),
prob=0.5),
ToFloat(inputs="x_aug2", outputs="x_aug2")
])
# step 2: prepare network
model_con, model_finetune = fe.build(model_fn=ResNet9, optimizer_fn=["adam", "adam"])
network = fe.Network(ops=[
LambdaOp(lambda x, y: tf.concat([x, y], axis=0), inputs=["x_aug", "x_aug2"], outputs="x_com"),
ModelOp(model=model_con, inputs="x_com", outputs="y_com"),
LambdaOp(lambda x: tf.split(x, 2, axis=0), inputs="y_com", outputs=["y_pred", "y_pred2"]),
NTXentOp(arg1="y_pred", arg2="y_pred2", outputs=["NTXent", "logit", "label"]),
UpdateOp(model=model_con, loss_name="NTXent")
])
# step 3: prepare estimator
traces = [
Accuracy(true_key="label", pred_key="logit", mode="train", output_name="contrastive_accuracy"),
ModelSaver(model=model_con, 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,
monitor_names="contrastive_accuracy")
estimator.fit()
return model_con, model_finetune
def test_saliency(self):
fe.estimator.enable_deterministic(200)
label_mapping = {
'airplane': 0,
'automobile': 1,
'bird': 2,
'cat': 3,
'deer': 4,
'dog': 5,
'frog': 6,
'horse': 7,
'ship': 8,
'truck': 9
}
batch_size = 32
train_data, eval_data = cifar10.load_data()
pipeline = fe.Pipeline(test_data=train_data,
batch_size=batch_size,
ops=[Normalize(inputs="x", outputs="x")],
num_process=0)
weight_path = os.path.abspath(
os.path.join(__file__, "..", "resources", "lenet_cifar10_tf.h5"))
model = fe.build(model_fn=lambda: LeNet(input_shape=(32, 32, 3)),
optimizer_fn="adam",
weights_path=weight_path)
network = fe.Network(
ops=[ModelOp(model=model, inputs="x", outputs="y_pred")])
save_dir = tempfile.mkdtemp()
traces = [
Saliency(model=model,
model_inputs="x",
class_key="y",
model_outputs="y_pred",
samples=5,
label_mapping=label_mapping),
ImageSaver(inputs="saliency", save_dir=save_dir)
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=5,
traces=traces,
log_steps=1000)
estimator.test()
ans_img_path = os.path.abspath(
os.path.join(__file__, "..", "resources", "saliency_figure.png"))
ans_img = img_to_rgb_array(ans_img_path)
output_img_path = os.path.join(save_dir, "saliency_test_epoch_5.png")
output_img = img_to_rgb_array(output_img_path)
self.assertTrue(check_img_similar(output_img, ans_img))
