def build_estimator(self):
"""
Define the estimator to run the experiment.
This will persist throughout the lifetime of the TaskRunner.
Args:
None
Returns:
estimator: Estimator object
"""
max_train_steps_per_epoch = None
max_eval_steps_per_epoch = None
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"),
]
estimator = fe.Estimator(
pipeline=self.data_loader.pipeline,
network=self.network,
epochs=2,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch,
max_eval_steps_per_epoch=max_eval_steps_per_epoch,
monitor_names=["base_ce", "adv_ce"],
log_steps=1000)
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 get_estimator(epsilon=0.04,
epochs=10,
batch_size=32,
max_train_steps_per_epoch=None,
max_eval_steps_per_epoch=None,
save_dir=tempfile.mkdtemp()):
# step 1
train_data, eval_data = cifair10.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)),
ChannelTranspose(inputs="x", outputs="x"),
])
# step 2
model = fe.build(model_fn=lambda: LeNet(input_shape=(3, 32, 32)),
optimizer_fn="adam")
network = fe.Network(ops=[
Watch(inputs="x"),
ModelOp(model=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=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=model, loss_name="avg_ce")
])
# step 3
traces = [
Accuracy(true_key="y", pred_key="y_pred", output_name="base accuracy"),
Accuracy(true_key="y",
pred_key="y_pred_adv",
output_name="adversarial accuracy"),
BestModelSaver(model=model,
save_dir=save_dir,
metric="adv_ce",
save_best_mode="min"),
]
estimator = fe.Estimator(
pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch,
max_eval_steps_per_epoch=max_eval_steps_per_epoch,
monitor_names=["base_ce", "adv_ce"])
return estimator
def get_estimator(epsilon=0.04,
epochs=20,
batch_size=32,
code_length=16,
train_steps_per_epoch=None,
eval_steps_per_epoch=None,
save_dir=tempfile.mkdtemp()):
# step 1
train_data, eval_data = cifair10.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)),
ChannelTranspose(inputs="x", outputs="x"),
Hadamard(inputs="y", outputs="y_code", n_classes=10)
],
num_process=0)
# step 2
model = fe.build(model_fn=lambda: EccLeNet(code_length=code_length), optimizer_fn="adam")
network = fe.Network(ops=[
Watch(inputs="x", mode=('eval', 'test')),
ModelOp(model=model, inputs="x", outputs="y_pred_code"),
Hinge(inputs=("y_pred_code", "y_code"), outputs="base_hinge"),
UpdateOp(model=model, loss_name="base_hinge"),
UnHadamard(inputs="y_pred_code", outputs="y_pred", n_classes=10, mode=('eval', 'test')),
# The adversarial attack:
FGSM(data="x", loss="base_hinge", outputs="x_adverse_hinge", epsilon=epsilon, mode=('eval', 'test')),
ModelOp(model=model, inputs="x_adverse_hinge", outputs="y_pred_adv_hinge_code", mode=('eval', 'test')),
Hinge(inputs=("y_pred_adv_hinge_code", "y_code"), outputs="adv_hinge", mode=('eval', 'test')),
UnHadamard(inputs="y_pred_adv_hinge_code", outputs="y_pred_adv_hinge", n_classes=10, mode=('eval', 'test')),
])
# step 3
traces = [
Accuracy(true_key="y", pred_key="y_pred", output_name="base_accuracy"),
Accuracy(true_key="y", pred_key="y_pred_adv_hinge", output_name="adversarial_accuracy"),
BestModelSaver(model=model, save_dir=save_dir, metric="base_hinge", save_best_mode="min", load_best_final=True)
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
train_steps_per_epoch=train_steps_per_epoch,
eval_steps_per_epoch=eval_steps_per_epoch,
monitor_names=["adv_hinge"])
return estimator
def setUpClass(cls):
x = np.array([[1, 2], [3, 4]])
x_pred = np.array([[1, 5, 3], [2, 1, 0]])
x_1d = np.array([2.5])
x_pred_1d = np.array([1])
x_1d_logit = np.array([1])
x_pred_1d_logit = np.array([2.5])
cls.data = Data({'x': x, 'x_pred': x_pred})
cls.data_1d = Data({'x': x_1d, 'x_pred': x_pred_1d})
cls.data_1d_logit = Data({'x': x_1d_logit, 'x_pred': x_pred_1d_logit})
cls.accuracy = Accuracy(true_key='x', pred_key='x_pred')
cls.accuracy_logit = Accuracy(true_key='x',
pred_key='x_pred',
from_logits=True)
def test_torch_binary_class(self):
with self.subTest("from_logit=False"):
trace = Accuracy(true_key="label", pred_key="pred", output_name=self.acc_key, from_logits=False)
batch = {"label": torch.tensor([0, 1])}
prediction = {"pred": torch.tensor([[0.3], [0.6]])} # pred > 0.5 => class 1
sim = TraceRun(trace=trace, batch=batch, prediction=prediction)
sim.run_trace()
self.assertEqual(sim.data_on_epoch_end[self.acc_key], 1.0)
with self.subTest("from_logit=True"):
trace = Accuracy(true_key="label", pred_key="pred", output_name=self.acc_key, from_logits=True)
batch = {"label": torch.tensor([0, 1])}
prediction = {"pred": torch.tensor([[-1], [1]])} # 1 / 1 + exp(-pred) > 0.5 => class 1
sim = TraceRun(trace=trace, batch=batch, prediction=prediction)
sim.run_trace()
self.assertEqual(sim.data_on_epoch_end[self.acc_key], 1.0)
def test_tf_one_hot_label(self):
trace = Accuracy(true_key="label", pred_key="pred", output_name=self.acc_key)
batch = {"label": tf.constant([[1, 0, 0], [0, 1, 0]])} # one-hot
prediction = {"pred": tf.constant([[1, 2, 3], [0.2, 0.5, 0.3]])}
sim = TraceRun(trace=trace, batch=batch, prediction=prediction)
sim.run_trace()
self.assertEqual(sim.data_on_epoch_end[self.acc_key], 0.5)
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 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=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=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(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 test_torch_class_index_label(self):
trace = Accuracy(true_key="label", pred_key="pred", output_name=self.acc_key)
batch = {"label": torch.tensor([0, 1])} # class index
prediction = {"pred": torch.tensor([[1, 2, 3], [0.2, 0.5, 0.3]])}
sim = TraceRun(trace=trace, batch=batch, prediction=prediction)
sim.run_trace()
self.assertEqual(sim.data_on_epoch_end[self.acc_key], 0.5)
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 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_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=2,
batch_size=32,
max_train_steps_per_epoch=None,
max_eval_steps_per_epoch=None,
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")
],
num_process=0)
# step 2
model = fe.build(model_fn=LeNet, optimizer_fn="adam")
print([f"{idx}: {x.name}" for idx, x in enumerate(model.submodules)])
network = fe.Network(ops=[
Watch(inputs="x"),
ModelOp(model=model,
inputs="x",
outputs=["y_pred", "embedding"],
intermediate_layers='dense'),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
GradientOp(finals="embedding", inputs="x", outputs="grads"),
UpdateOp(model=model, loss_name="ce")
])
# step 3
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
Inspector(),
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)),
TensorBoard(log_dir="tf_logs",
write_embeddings="embedding",
embedding_labels="y")
]
estimator = fe.Estimator(
pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch,
max_eval_steps_per_epoch=max_eval_steps_per_epoch)
return estimator
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(level,
num_augment,
epochs=24,
batch_size=512,
max_train_steps_per_epoch=None,
max_eval_steps_per_epoch=None):
assert 0 <= level <= 10, "the level should be between 0 and 10"
train_data, test_data = load_data()
aug_ops = [
OneOf(
Rotate(level=level, inputs="x", outputs="x", mode="train"),
Identity(level=level, inputs="x", outputs="x", mode="train"),
AutoContrast(level=level, inputs="x", outputs="x", mode="train"),
Equalize(level=level, inputs="x", outputs="x", mode="train"),
Posterize(level=level, inputs="x", outputs="x", mode="train"),
Solarize(level=level, inputs="x", outputs="x", mode="train"),
Sharpness(level=level, inputs="x", outputs="x", mode="train"),
Contrast(level=level, inputs="x", outputs="x", mode="train"),
Color(level=level, inputs="x", outputs="x", mode="train"),
Brightness(level=level, inputs="x", outputs="x", mode="train"),
ShearX(level=level, inputs="x", outputs="x", mode="train"),
ShearY(level=level, inputs="x", outputs="x", mode="train"),
TranslateX(level=level, inputs="x", outputs="x", mode="train"),
TranslateY(level=level, inputs="x", outputs="x", mode="train"),
) for _ in range(num_augment)
]
pipeline = fe.Pipeline(train_data=train_data,
test_data=test_data,
batch_size=batch_size,
ops=aug_ops + [
Normalize(inputs="x",
outputs="x",
mean=(0.4914, 0.4822, 0.4465),
std=(0.2471, 0.2435, 0.2616)),
ChannelTranspose(inputs="x", outputs="x"),
])
model = fe.build(model_fn=MyModel, 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")
])
estimator = fe.Estimator(
pipeline=pipeline,
network=network,
epochs=epochs,
traces=Accuracy(true_key="y", pred_key="y_pred"),
max_train_steps_per_epoch=max_train_steps_per_epoch,
max_eval_steps_per_epoch=max_eval_steps_per_epoch)
return estimator
def get_estimator(max_len=20,
epochs=10,
batch_size=64,
max_train_steps_per_epoch=None,
max_eval_steps_per_epoch=None,
pretrained_model='bert-base-uncased',
save_dir=tempfile.mkdtemp(),
data_dir=None):
# step 1 prepare data
train_data, eval_data, data_vocab, label_vocab = german_ner.load_data(root_dir=data_dir)
tokenizer = BertTokenizer.from_pretrained(pretrained_model, do_lower_case=True)
tag2idx = char2idx(label_vocab)
pipeline = fe.Pipeline(
train_data=train_data,
eval_data=eval_data,
batch_size=batch_size,
ops=[
Tokenize(inputs="x", outputs="x", tokenize_fn=tokenizer.tokenize),
WordtoId(inputs="x", outputs="x", mapping=tokenizer.convert_tokens_to_ids),
WordtoId(inputs="y", outputs="y", mapping=tag2idx),
PadSequence(max_len=max_len, inputs="x", outputs="x"),
PadSequence(max_len=max_len, value=len(tag2idx), inputs="y", outputs="y"),
AttentionMask(inputs="x", outputs="x_masks")
])
# step 2. prepare model
bert_config = BertConfig.from_pretrained(pretrained_model)
num_hidden_layers = bert_config.to_dict()['num_hidden_layers']
head_masks = [None] * num_hidden_layers
model = fe.build(model_fn=lambda: NERModel(head_masks=head_masks, pretrained_model=pretrained_model),
optimizer_fn=lambda x: torch.optim.Adam(x, lr=1e-5))
network = fe.Network(ops=[
ModelOp(model=model, inputs=["x", "x_masks"], outputs="y_pred"),
Reshape(inputs="y", outputs="y", shape=(-1, )),
Reshape(inputs="y_pred", outputs="y_pred", shape=(-1, 24)),
CrossEntropy(inputs=("y_pred", "y"), outputs="loss"),
UpdateOp(model=model, loss_name="loss")
])
traces = [Accuracy(true_key="y", pred_key="y_pred"), BestModelSaver(model=model, save_dir=save_dir)]
# step 3 prepare estimator
estimator = fe.Estimator(network=network,
pipeline=pipeline,
epochs=epochs,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch,
max_eval_steps_per_epoch=max_eval_steps_per_epoch)
return estimator
def search_max_lr(pipeline, model, network, epochs):
traces = [
Accuracy(true_key="y", pred_key="y_pred"), LRScheduler(model=model, lr_fn=lambda step: linear_increase(step))
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
max_train_steps_per_epoch=10,
log_steps=10)
print("Running LR range test for super convergence. It will take a while...")
with Suppressor():
summary = estimator.fit("LR_range_test")
best_step = max(summary.history["eval"]["accuracy"].items(), key=lambda k: k[1])[0]
max_lr = summary.history["train"]["model_lr"][best_step]
return max_lr
def get_estimator(epochs=20,
batch_size=32,
max_train_steps_per_epoch=None,
max_eval_steps_per_epoch=None,
save_dir=tempfile.mkdtemp()):
# step 1. prepare data
train_data, eval_data = breast_cancer.load_data()
test_data = eval_data.split(0.5)
# Apply some global pre-processing to the data
scaler = StandardScaler()
train_data["x"] = scaler.fit_transform(train_data["x"])
eval_data["x"] = scaler.transform(eval_data["x"])
test_data["x"] = scaler.transform(test_data["x"])
pipeline = fe.Pipeline(train_data=train_data,
eval_data=eval_data,
test_data=test_data,
batch_size=batch_size)
# step 2. prepare model
model = fe.build(model_fn=create_dnn, optimizer_fn="adam")
network = fe.Network(ops=[
ModelOp(inputs="x", model=model, 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")
]
estimator = fe.Estimator(
pipeline=pipeline,
network=network,
epochs=epochs,
log_steps=10,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch,
max_eval_steps_per_epoch=max_eval_steps_per_epoch)
return estimator
def setUpClass(cls):
cls.data_np = Data({'loss': np.NaN})
cls.data_tf = Data({'loss': tf.constant(np.NaN)})
cls.data_torch = Data({'loss': torch.tensor(np.NaN)})
cls.expected_msg = "FastEstimator-TerminateOnNaN: NaN Detected in: loss"
cls.expected_loss_keys = {"ce"}
cls.expected_all_keys = {"ce", "accuracy", "f1_score"}
tf_model = fe.build(model_fn=one_layer_tf_model, optimizer_fn='adam')
cls.network = fe.Network(ops=[
ModelOp(model=tf_model, inputs="x", outputs="y"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
UpdateOp(model=tf_model, loss_name="ce")
])
cls.traces = [
Accuracy(true_key="y", pred_key="y_pred", output_name="accuracy"),
F1Score(true_key="y", pred_key="y_pred", output_name="f1_score")
]
def finetune(weights_path,
batch_size,
epochs,
model_dir=tempfile.mkdtemp(),
train_steps_per_epoch=None,
eval_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", 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)
])
_, model = fe.build(
model_fn=lambda: vision_transformer(num_class=10,
weights_path=weights_path,
image_size=(32, 32, 3),
patch_size=4,
num_layers=6,
em_dim=256,
num_heads=8,
dff=512),
optimizer_fn=[None, lambda: tf.optimizers.SGD(0.01, momentum=0.9)])
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce", from_logits=True),
UpdateOp(model=model, loss_name="ce")
])
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
BestModelSaver(model=model, save_dir=model_dir, metric="accuracy", save_best_mode="max")
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
train_steps_per_epoch=train_steps_per_epoch,
eval_steps_per_epoch=eval_steps_per_epoch)
estimator.fit(warmup=False)
def get_estimator(epochs=24, batch_size=512, 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,
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)),
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=FastCifar, 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 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=lr_schedule)
]
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(batch_size=8,
epochs=50,
train_steps_per_epoch=None,
eval_steps_per_epoch=None,
save_dir=tempfile.mkdtemp(),
data_dir=None):
# load CUB200 dataset.
train_data = cub200.load_data(root_dir=data_dir)
eval_data = train_data.split(0.3)
test_data = eval_data.split(0.5)
# step 1, pipeline
pipeline = fe.Pipeline(batch_size=batch_size,
train_data=train_data,
eval_data=eval_data,
test_data=test_data,
ops=[
ReadImage(inputs="image",
outputs="image",
parent_path=train_data.parent_path),
Normalize(inputs="image",
outputs="image",
mean=1.0,
std=1.0,
max_pixel_value=127.5),
ReadMat(file='annotation',
keys="seg",
parent_path=train_data.parent_path),
LongestMaxSize(max_size=512,
image_in="image",
image_out="image",
mask_in="seg",
mask_out="seg"),
PadIfNeeded(min_height=512,
min_width=512,
image_in="image",
image_out="image",
mask_in="seg",
mask_out="seg",
border_mode=cv2.BORDER_CONSTANT,
value=0,
mask_value=0),
ShiftScaleRotate(
image_in="image",
mask_in="seg",
image_out="image",
mask_out="seg",
mode="train",
shift_limit=0.2,
rotate_limit=15.0,
scale_limit=0.2,
border_mode=cv2.BORDER_CONSTANT,
value=0,
mask_value=0),
Sometimes(
HorizontalFlip(image_in="image",
mask_in="seg",
image_out="image",
mask_out="seg",
mode="train")),
Reshape(shape=(512, 512, 1),
inputs="seg",
outputs="seg")
])
# step 2, network
resunet50 = fe.build(model_fn=ResUnet50,
model_name="resunet50",
optimizer_fn=lambda: tf.optimizers.Adam(1e-4))
uncertainty = fe.build(model_fn=UncertaintyLossNet,
model_name="uncertainty",
optimizer_fn=lambda: tf.optimizers.Adam(2e-5))
network = fe.Network(ops=[
ModelOp(inputs='image',
model=resunet50,
outputs=["label_pred", "mask_pred"]),
CrossEntropy(inputs=["label_pred", "label"],
outputs="cls_loss",
form="sparse",
average_loss=False),
CrossEntropy(inputs=["mask_pred", "seg"],
outputs="seg_loss",
form="binary",
average_loss=False),
ModelOp(inputs=["cls_loss", "seg_loss"],
model=uncertainty,
outputs="total_loss"),
ReduceLoss(inputs="total_loss", outputs="total_loss"),
UpdateOp(model=resunet50, loss_name="total_loss"),
UpdateOp(model=uncertainty, loss_name="total_loss")
])
# step 3, estimator
traces = [
Accuracy(true_key="label", pred_key="label_pred"),
Dice(true_key="seg", pred_key='mask_pred'),
BestModelSaver(model=resunet50,
save_dir=save_dir,
metric="total_loss",
save_best_mode="min"),
LRScheduler(model=resunet50,
lr_fn=lambda step: cosine_decay(
step, cycle_length=26400, init_lr=1e-4))
]
estimator = fe.Estimator(network=network,
pipeline=pipeline,
traces=traces,
epochs=epochs,
train_steps_per_epoch=train_steps_per_epoch,
eval_steps_per_epoch=eval_steps_per_epoch,
log_steps=500)
return estimator
def get_estimator(epochs=150,
batch_size=32,
save_dir=tempfile.mkdtemp(),
train_steps_per_epoch=None,
eval_steps_per_epoch=None):
# step 1: prepare dataset
train_data, eval_data = load_data()
pipeline = fe.Pipeline(train_data=train_data,
eval_data=eval_data,
batch_size=batch_size * get_num_devices(),
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)
])
# step 2: prepare network
model = fe.build(
model_fn=lambda: pyramidnet_cifar(inputs_shape=(32, 32, 3),
depth=272,
alpha=200,
num_classes=10,
bottleneck=True),
optimizer_fn=lambda: tfa.optimizers.SGDW(
weight_decay=0.0001, lr=0.1, momentum=0.9))
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce", from_logits=True),
UpdateOp(model=model, loss_name="ce")
])
# step 3 prepare estimator
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
LRScheduler(model=model, lr_fn=lr_schedule),
BestModelSaver(model=model,
save_dir=save_dir,
metric="accuracy",
save_best_mode="max")
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
train_steps_per_epoch=train_steps_per_epoch,
eval_steps_per_epoch=eval_steps_per_epoch)
return estimator
def finetune(pretrained_model,
batch_size,
epochs,
model_dir=tempfile.mkdtemp(),
train_steps_per_epoch=None,
eval_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",
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")
])
model = fe.build(model_fn=lambda: ViTModel(num_classes=100,
image_size=32,
patch_size=4,
num_layers=6,
num_channels=3,
em_dim=256,
num_heads=8,
ff_dim=512),
optimizer_fn=lambda x: torch.optim.SGD(
x, lr=0.01, momentum=0.9, weight_decay=1e-4))
# load the encoder's weight
if hasattr(model, "module"):
model.module.vit_encoder.load_state_dict(
pretrained_model.module.vit_encoder.state_dict())
else:
model.vit_encoder.load_state_dict(
pretrained_model.vit_encoder.state_dict())
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce", from_logits=True),
UpdateOp(model=model, loss_name="ce")
])
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
BestModelSaver(model=model,
save_dir=model_dir,
metric="accuracy",
save_best_mode="max")
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
train_steps_per_epoch=train_steps_per_epoch,
eval_steps_per_epoch=eval_steps_per_epoch)
estimator.fit(warmup=False)
def pretrain_model(epochs, batch_size, train_steps_per_epoch, save_dir):
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",
mode="train"),
# 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),
ChannelTranspose(inputs="x_aug", outputs="x_aug"),
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),
ChannelTranspose(inputs="x_aug2", outputs="x_aug2"),
ToFloat(inputs="x_aug2", outputs="x_aug2")
])
model_con = fe.build(model_fn=lambda: ResNet9OneLayerHead(length=128),
optimizer_fn="adam")
network = fe.Network(ops=[
LambdaOp(lambda x, y: torch.cat([x, y], dim=0),
inputs=["x_aug", "x_aug2"],
outputs="x_com"),
ModelOp(model=model_con, inputs="x_com", outputs="y_com"),
LambdaOp(lambda x: torch.chunk(x, 2, dim=0),
inputs="y_com",
outputs=["y_pred", "y_pred2"],
mode="train"),
NTXentOp(arg1="y_pred",
arg2="y_pred2",
outputs=["NTXent", "logit", "label"],
mode="train"),
UpdateOp(model=model_con, loss_name="NTXent")
])
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,
train_steps_per_epoch=train_steps_per_epoch)
estimator.fit()
return model_con
def get_estimator(epochs=200,
batch_size=128,
max_train_steps_per_epoch=None,
max_eval_steps_per_epoch=None,
save_dir=tempfile.mkdtemp(),
data_dir=None):
# step 1. prepare pipeline
train_data, eval_data = omniglot.load_data(root_dir=data_dir)
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=[
ReadImage(inputs="x_a",
outputs="x_a",
color_flag='gray'),
ReadImage(inputs="x_b",
outputs="x_b",
color_flag='gray'),
Sometimes(ShiftScaleRotate(image_in="x_a",
image_out="x_a",
shift_limit=0.05,
scale_limit=0.2,
rotate_limit=10,
mode="train"),
prob=0.89),
Sometimes(ShiftScaleRotate(image_in="x_b",
image_out="x_b",
shift_limit=0.05,
scale_limit=0.2,
rotate_limit=10,
mode="train"),
prob=0.89),
Minmax(inputs="x_a", outputs="x_a"),
Minmax(inputs="x_b", outputs="x_b")
])
# step 2. prepare model
model = fe.build(model_fn=siamese_network,
model_name="siamese_net",
optimizer_fn="adam")
network = fe.Network(ops=[
ModelOp(inputs=["x_a", "x_b"], model=model, outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="loss", form="binary"),
UpdateOp(model=model, loss_name="loss")
])
# step 3.prepare estimator
traces = [
LRScheduler(model=model, lr_fn=lr_schedule),
Accuracy(true_key="y", pred_key="y_pred"),
OneShotAccuracy(dataset=eval_data,
model=model,
output_name='one_shot_accuracy'),
BestModelSaver(model=model,
save_dir=save_dir,
metric="one_shot_accuracy",
save_best_mode="max"),
EarlyStopping(monitor="one_shot_accuracy",
patience=20,
compare='max',
mode="eval")
]
estimator = fe.Estimator(
network=network,
pipeline=pipeline,
epochs=epochs,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch,
max_eval_steps_per_epoch=max_eval_steps_per_epoch)
return estimator
