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