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 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 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(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(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 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 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 = cifair10.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))
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(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)
from os import pipe
import fastestimator as fe
from fastestimator.dataset.data.cifair10 import load_data
from fastestimator.op.numpyop.meta import Sometimes
from fastestimator.op.numpyop.multivariate import HorizontalFlip, PadIfNeeded, RandomCrop
from fastestimator.op.numpyop.univariate import CoarseDropout, Normalize, Onehot
train_data, test_data = load_data()
pipeline = fe.Pipeline(train_data=train_data,
test_data=test_data,
batch_size={
"train": 128,
"test": 32
},
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(
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=12, batch_size=512, save_dir=tempfile.mkdtemp()):
# epoch 1-10, train on cifair100, epoch 11-end: train on cifar10
cifair10_train, cifari10_test = cifair10.load_data()
cifair100_train, _ = cifair100.load_data()
train_ds = EpochScheduler({1: cifair100_train, 11: cifair10_train})
pipeline = fe.Pipeline(train_data=train_ds,
test_data=cifari10_test,
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")
])
# step 2: prepare network
backbone = fe.build(model_fn=lambda: Backbone(input_size=(3, 32, 32)),
optimizer_fn="adam")
cls_head_cifar100 = fe.build(model_fn=lambda: Classifier(classes=100),
optimizer_fn="adam")
cls_head_cifar10 = fe.build(model_fn=lambda: Classifier(classes=10),
optimizer_fn="adam")
# if you want to save the final cifar10 model, you can build a model then provide it to ModelSaver
# final_model_cifar10 = fe.build(model_fn=lambda: MyModel(backbone, cls_head_cifar10), optimizer_fn=None)
# epoch 1-10: train backbone and cls_head_cifar100, epoch 11-end: train cls_head_cifar10 only
ModelOp_cls_head = EpochScheduler({
1:
ModelOp(model=cls_head_cifar100, inputs="feature", outputs="y_pred"),
11:
ModelOp(model=cls_head_cifar10, inputs="feature", outputs="y_pred"),
})
UpdateOp_backbone = EpochScheduler({
1:
UpdateOp(model=backbone, loss_name="ce"),
11:
None
})
UpdateOp_cls_head = EpochScheduler({
1:
UpdateOp(model=cls_head_cifar100, loss_name="ce"),
11:
UpdateOp(model=cls_head_cifar10, loss_name="ce")
})
network = fe.Network(ops=[
ModelOp(model=backbone, inputs="x",
outputs="feature"), ModelOp_cls_head,
CrossEntropy(inputs=("y_pred", "y"), outputs="ce", from_logits=True),
UpdateOp_backbone, UpdateOp_cls_head
])
traces = [Accuracy(true_key="y", pred_key="y_pred")]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces)
return estimator
