def test_tf_model(self): def update(): with tf.GradientTape(persistent=True) as tape: self.state['tape'] = tape pred = fe.backend.feed_forward(self.tf_model, self.tf_input_data) loss = fe.backend.mean_squared_error(y_pred=pred, y_true=self.tf_y) op.forward(data=loss, state=self.state) op = UpdateOp(model=self.tf_model, loss_name='loss') op.build("tf") strategy = tf.distribute.get_strategy() if isinstance(strategy, tf.distribute.MirroredStrategy): strategy.run(update, args=()) else: update() bms = BestModelSaver(model=self.tf_model, save_dir=self.save_dir) bms.on_epoch_end(data=self.data) m2 = fe.build(model_fn=one_layer_model_without_weights, optimizer_fn='adam') fe.backend.load_model(m2, os.path.join(self.save_dir, 'tf_best_loss.h5')) self.assertTrue( is_equal(m2.trainable_variables, self.tf_model.trainable_variables))
def test_torch_model(self): op = UpdateOp(model=self.torch_model, loss_name='loss') pred = fe.backend.feed_forward(self.torch_model, self.torch_input_data) loss = fe.backend.mean_squared_error(y_pred=pred, y_true=self.torch_y) output = op.forward(data=loss, state=self.state) bms = BestModelSaver(model=self.torch_model, save_dir=self.save_dir) bms.on_epoch_end(data=self.data) m2 = fe.build(model_fn=MultiLayerTorchModelWithoutWeights, optimizer_fn='adam') fe.backend.load_model( m2, os.path.join(self.save_dir, 'torch_best_loss.pt')) self.assertTrue( is_equal(list(m2.parameters()), list(self.torch_model.parameters())))
def test_tf_model(self): op = UpdateOp(model=self.tf_model, loss_name='loss') with tf.GradientTape(persistent=True) as tape: self.state['tape'] = tape pred = fe.backend.feed_forward(self.tf_model, self.tf_input_data) loss = fe.backend.mean_squared_error(y_pred=pred, y_true=self.tf_y) output = op.forward(data=loss, state=self.state) bms = BestModelSaver(model=self.tf_model, save_dir=self.save_dir) bms.on_epoch_end(data=self.data) m2 = fe.build(model_fn=one_layer_model_without_weights, optimizer_fn='adam') fe.backend.load_model(m2, os.path.join(self.save_dir, 'tf_best_loss.h5')) self.assertTrue( is_equal(m2.trainable_variables, self.tf_model.trainable_variables))
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 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(batch_size=100, epochs=20, max_train_steps_per_epoch=None, save_dir=tempfile.mkdtemp()): train_data, _ = load_data() pipeline = fe.Pipeline(train_data=train_data, batch_size=batch_size, ops=[ ExpandDims(inputs="x", outputs="x", axis=0), Minmax(inputs="x", outputs="x"), Binarize(inputs="x", outputs="x", threshold=0.5), ]) encode_model = fe.build(model_fn=EncoderNet, optimizer_fn="adam", model_name="encoder") decode_model = fe.build(model_fn=DecoderNet, optimizer_fn="adam", model_name="decoder") network = fe.Network(ops=[ ModelOp(model=encode_model, inputs="x", outputs="meanlogvar"), SplitOp(inputs="meanlogvar", outputs=("mean", "logvar")), ReparameterizepOp(inputs=("mean", "logvar"), outputs="z"), ModelOp(model=decode_model, inputs="z", outputs="x_logit"), CrossEntropy(inputs=("x_logit", "x"), outputs="cross_entropy"), CVAELoss(inputs=("cross_entropy", "mean", "logvar", "z"), outputs="loss"), UpdateOp(model=encode_model, loss_name="loss"), UpdateOp(model=decode_model, loss_name="loss"), ]) traces = [ BestModelSaver(model=encode_model, save_dir=save_dir), BestModelSaver(model=decode_model, save_dir=save_dir) ] estimator = fe.Estimator( pipeline=pipeline, network=network, epochs=epochs, traces=traces, max_train_steps_per_epoch=max_train_steps_per_epoch) return estimator
def get_estimator(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=50, batch_size=128, max_train_steps_per_epoch=None, max_eval_steps_per_epoch=None, save_dir=tempfile.mkdtemp()): # step 1 train_data, eval_data = cifair100.load_data() # Add label noise to simulate real-world labeling problems corrupt_dataset(train_data) test_data = eval_data.split(range(len(eval_data) // 2)) 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)), 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 model = fe.build(model_fn=big_lenet, optimizer_fn='adam') network = fe.Network(ops=[ ModelOp(model=model, inputs="x", outputs="y_pred"), SuperLoss(CrossEntropy(inputs=("y_pred", "y"), outputs="ce"), output_confidence="confidence"), UpdateOp(model=model, loss_name="ce") ]) # step 3 traces = [ MCC(true_key="y", pred_key="y_pred"), BestModelSaver(model=model, save_dir=save_dir, metric="mcc", save_best_mode="max", load_best_final=True), LabelTracker(metric="confidence", label="data_labels", label_mapping={ "Normal": 0, "Corrupted": 1 }, mode="train", outputs="label_confidence"), ImageSaver(inputs="label_confidence", save_dir=save_dir, mode="train"), ] 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=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(data_dir=None, model_dir=tempfile.mkdtemp(), epochs=20, em_dim=128, batch_size=32, max_train_steps_per_epoch=None, max_eval_steps_per_epoch=None): train_ds, eval_ds, test_ds = tednmt.load_data(data_dir, translate_option="pt_to_en") pt_tokenizer = BertTokenizer.from_pretrained("neuralmind/bert-base-portuguese-cased") en_tokenizer = BertTokenizer.from_pretrained("bert-base-uncased") pipeline = fe.Pipeline( train_data=train_ds, eval_data=eval_ds, test_data=test_ds, batch_size=batch_size, ops=[ Encode(inputs="source", outputs="source", tokenizer=pt_tokenizer), Encode(inputs="target", outputs="target", tokenizer=en_tokenizer) ], pad_value=0) model = fe.build( model_fn=lambda: Transformer(num_layers=4, em_dim=em_dim, num_heads=8, ff_dim=512, input_vocab=pt_tokenizer.vocab_size, target_vocab=en_tokenizer.vocab_size, max_pos_enc=1000, max_pos_dec=1000), optimizer_fn="adam") network = fe.Network(ops=[ ShiftData(inputs="target", outputs=("target_inp", "target_real")), CreateMasks(inputs=("source", "target_inp"), outputs=("encode_pad_mask", "decode_pad_mask", "dec_look_ahead_mask")), ModelOp(model=model, inputs=("source", "target_inp", "encode_pad_mask", "decode_pad_mask", "dec_look_ahead_mask"), outputs="pred"), MaskedCrossEntropy(inputs=("pred", "target_real"), outputs="ce"), UpdateOp(model=model, loss_name="ce") ]) traces = [ MaskedAccuracy(inputs=("pred", "target_real"), outputs="masked_acc", mode="!train"), BestModelSaver(model=model, save_dir=model_dir, metric="masked_acc", save_best_mode="max"), LRScheduler(model=model, lr_fn=lambda step: lr_fn(step, em_dim)) ] estimator = fe.Estimator(pipeline=pipeline, network=network, traces=traces, epochs=epochs, 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 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(epochs=30, batch_size=128, seq_length=20, vocab_size=10000, data_dir=None, train_steps_per_epoch=None, save_dir=tempfile.mkdtemp()): train_data, eval_data, _, _ = load_data(root_dir=data_dir, seq_length=seq_length + 1) pipeline = fe.Pipeline(train_data=train_data, eval_data=eval_data, ops=[ CreateInputAndTarget(inputs="x", outputs=("x", "y")), Batch(batch_size=batch_size, drop_last=True) ]) # step 2 model = fe.build( model_fn=lambda: BuildModel( vocab_size, embedding_dim=300, rnn_units=600), optimizer_fn=lambda x: torch.optim.SGD(x, lr=1.0, momentum=0.9)) network = fe.Network(ops=[ DimesionAdjust(inputs=("x", "y"), outputs=("x", "y")), ModelOp(model=model, inputs="x", outputs="y_pred", mode=None), CrossEntropy(inputs=("y_pred", "y"), outputs="ce", form="sparse", from_logits=True), UpdateOp(model=model, loss_name="ce") ]) # step 3 traces = [ Perplexity(inputs="ce", outputs="perplexity", mode="eval"), LRScheduler(model=model, lr_fn=lambda step: lr_schedule(step, init_lr=1.0)), BestModelSaver(model=model, save_dir=save_dir, metric='perplexity', save_best_mode='min', load_best_final=True), EarlyStopping(monitor="perplexity", patience=5) ] 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=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 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 get_estimator(data_dir=None, model_dir=tempfile.mkdtemp(), batch_size=16, epochs=13, max_train_steps_per_epoch=None, max_eval_steps_per_epoch=None, image_size=512, num_classes=90): # pipeline train_ds, eval_ds = mscoco.load_data(root_dir=data_dir) pipeline = fe.Pipeline( train_data=train_ds, eval_data=eval_ds, batch_size=batch_size, ops=[ ReadImage(inputs="image", outputs="image"), LongestMaxSize(image_size, image_in="image", image_out="image", bbox_in="bbox", bbox_out="bbox", bbox_params=BboxParams("coco", min_area=1.0)), PadIfNeeded( image_size, image_size, border_mode=cv2.BORDER_CONSTANT, image_in="image", image_out="image", bbox_in="bbox", bbox_out="bbox", bbox_params=BboxParams("coco", min_area=1.0), ), Sometimes( HorizontalFlip(mode="train", image_in="image", image_out="image", bbox_in="bbox", bbox_out="bbox", bbox_params='coco')), # normalize from uint8 to [-1, 1] Normalize(inputs="image", outputs="image", mean=1.0, std=1.0, max_pixel_value=127.5), ShiftLabel(inputs="bbox", outputs="bbox"), AnchorBox(inputs="bbox", outputs="anchorbox", width=image_size, height=image_size), ChannelTranspose(inputs="image", outputs="image") ], pad_value=0) # network model = fe.build(model_fn=lambda: RetinaNet(num_classes=num_classes), optimizer_fn=lambda x: torch.optim.SGD( x, lr=2e-4, momentum=0.9, weight_decay=0.0001)) network = fe.Network(ops=[ ModelOp(model=model, inputs="image", outputs=["cls_pred", "loc_pred"]), RetinaLoss(inputs=["anchorbox", "cls_pred", "loc_pred"], outputs=["total_loss", "focal_loss", "l1_loss"]), UpdateOp(model=model, loss_name="total_loss"), PredictBox(input_shape=(image_size, image_size, 3), inputs=["cls_pred", "loc_pred"], outputs="pred", mode="eval") ]) # estimator traces = [ LRScheduler(model=model, lr_fn=lr_fn), BestModelSaver(model=model, save_dir=model_dir, metric='mAP', save_best_mode="max"), MeanAveragePrecision(num_classes=num_classes, true_key='bbox', pred_key='pred', mode="eval") ] 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=["l1_loss", "focal_loss"]) return estimator
def get_estimator(epochs=20, batch_size=128, max_train_steps_per_epoch=None, save_dir=tempfile.mkdtemp()): # Dataset Creation (x_train, y_train), (x_eval, y_eval) = tf.keras.datasets.mnist.load_data() x_eval0, y_eval0 = x_eval[np.where((y_eval == 1))], np.ones(y_eval[np.where((y_eval == 1))].shape) x_eval1, y_eval1 = x_eval[np.where((y_eval != 1))], y_eval[np.where((y_eval != 1))] # Ensuring outliers comprise 50% of the dataset index = np.random.choice(x_eval1.shape[0], int(x_eval0.shape[0]), replace=False) x_eval1, y_eval1 = x_eval1[index], np.zeros(y_eval1[index].shape) x_train, y_train = x_train[np.where((y_train == 1))], np.zeros(y_train[np.where((y_train == 1))].shape) train_data = fe.dataset.NumpyDataset({"x": x_train, "y": y_train}) x_eval, y_eval = np.concatenate([x_eval0, x_eval1]), np.concatenate([y_eval0, y_eval1]) eval_data = fe.dataset.NumpyDataset({"x": x_eval, "y": y_eval}) pipeline = fe.Pipeline( train_data=train_data, eval_data=eval_data, batch_size=batch_size, ops=[ ExpandDims(inputs="x", outputs="x"), Normalize(inputs="x", outputs="x", mean=1.0, std=1.0, max_pixel_value=127.5), LambdaOp(fn=lambda x: x + np.random.normal(loc=0.0, scale=0.155, size=(28, 28, 1)), inputs="x", outputs="x_w_noise", mode="train") ]) recon_model = fe.build(model_fn=reconstructor, optimizer_fn=lambda: tf.optimizers.RMSprop(2e-4), model_name="reconstructor") disc_model = fe.build(model_fn=discriminator, optimizer_fn=lambda: tf.optimizers.RMSprop(1e-4), model_name="discriminator") network = fe.Network(ops=[ ModelOp(model=recon_model, inputs="x_w_noise", outputs="x_fake", mode="train"), ModelOp(model=recon_model, inputs="x", outputs="x_fake", mode="eval"), ModelOp(model=disc_model, inputs="x_fake", outputs="fake_score"), ModelOp(model=disc_model, inputs="x", outputs="true_score"), RLoss(inputs=("fake_score", "x_fake", "x"), outputs="rloss"), UpdateOp(model=recon_model, loss_name="rloss"), DLoss(inputs=("true_score", "fake_score"), outputs="dloss"), UpdateOp(model=disc_model, loss_name="dloss"), ]) traces = [ F1AUCScores(true_key="y", pred_key="fake_score", mode="eval", output_name=["auc_score", "f1_score"]), BestModelSaver(model=recon_model, save_dir=save_dir, metric='f1_score', save_best_mode='max'), BestModelSaver(model=disc_model, save_dir=save_dir, metric='f1_score', save_best_mode='max'), ] estimator = fe.Estimator(pipeline=pipeline, network=network, epochs=epochs, traces=traces, max_train_steps_per_epoch=max_train_steps_per_epoch, log_steps=50) return estimator
def get_estimator(data_dir=None, model_dir=tempfile.mkdtemp(), epochs=200, batch_size_per_gpu=32, train_steps_per_epoch=None, eval_steps_per_epoch=None): num_device = get_num_devices() train_ds, val_ds = mscoco.load_data(root_dir=data_dir) train_ds = PreMosaicDataset(mscoco_ds=train_ds) batch_size = num_device * batch_size_per_gpu pipeline = fe.Pipeline( train_data=train_ds, eval_data=val_ds, ops=[ ReadImage(inputs=("image1", "image2", "image3", "image4"), outputs=("image1", "image2", "image3", "image4"), mode="train"), ReadImage(inputs="image", outputs="image", mode="eval"), LongestMaxSize(max_size=640, image_in="image1", bbox_in="bbox1", bbox_params=BboxParams("coco", min_area=1.0), mode="train"), LongestMaxSize(max_size=640, image_in="image2", bbox_in="bbox2", bbox_params=BboxParams("coco", min_area=1.0), mode="train"), LongestMaxSize(max_size=640, image_in="image3", bbox_in="bbox3", bbox_params=BboxParams("coco", min_area=1.0), mode="train"), LongestMaxSize(max_size=640, image_in="image4", bbox_in="bbox4", bbox_params=BboxParams("coco", min_area=1.0), mode="train"), LongestMaxSize(max_size=640, image_in="image", bbox_in="bbox", bbox_params=BboxParams("coco", min_area=1.0), mode="eval"), PadIfNeeded(min_height=640, min_width=640, image_in="image", bbox_in="bbox", bbox_params=BboxParams("coco", min_area=1.0), mode="eval", border_mode=cv2.BORDER_CONSTANT, value=(114, 114, 114)), CombineMosaic(inputs=("image1", "image2", "image3", "image4", "bbox1", "bbox2", "bbox3", "bbox4"), outputs=("image", "bbox"), mode="train"), CenterCrop(height=640, width=640, image_in="image", bbox_in="bbox", bbox_params=BboxParams("coco", min_area=1.0), mode="train"), Sometimes( HorizontalFlip(image_in="image", bbox_in="bbox", bbox_params=BboxParams("coco", min_area=1.0), mode="train")), HSVAugment(inputs="image", outputs="image", mode="train"), ToArray(inputs="bbox", outputs="bbox", dtype="float32"), CategoryID2ClassID(inputs="bbox", outputs="bbox"), GTBox(inputs="bbox", outputs=("gt_sbbox", "gt_mbbox", "gt_lbbox"), image_size=640), Delete(keys=("image1", "image2", "image3", "image4", "bbox1", "bbox2", "bbox3", "bbox4", "bbox"), mode="train"), Delete(keys="image_id", mode="eval"), Batch(batch_size=batch_size, pad_value=0) ]) init_lr = 1e-2 / 64 * batch_size model = fe.build( lambda: YoloV5(w=640, h=640, c=3), optimizer_fn=lambda x: torch.optim.SGD( x, lr=init_lr, momentum=0.937, weight_decay=0.0005, nesterov=True), mixed_precision=True) network = fe.Network(ops=[ RescaleTranspose(inputs="image", outputs="image"), ModelOp(model=model, inputs="image", outputs=("pred_s", "pred_m", "pred_l")), DecodePred(inputs=("pred_s", "pred_m", "pred_l"), outputs=("pred_s", "pred_m", "pred_l")), ComputeLoss(inputs=("pred_s", "gt_sbbox"), outputs=("sbbox_loss", "sconf_loss", "scls_loss")), ComputeLoss(inputs=("pred_m", "gt_mbbox"), outputs=("mbbox_loss", "mconf_loss", "mcls_loss")), ComputeLoss(inputs=("pred_l", "gt_lbbox"), outputs=("lbbox_loss", "lconf_loss", "lcls_loss")), Average(inputs=("sbbox_loss", "mbbox_loss", "lbbox_loss"), outputs="bbox_loss"), Average(inputs=("sconf_loss", "mconf_loss", "lconf_loss"), outputs="conf_loss"), Average(inputs=("scls_loss", "mcls_loss", "lcls_loss"), outputs="cls_loss"), Average(inputs=("bbox_loss", "conf_loss", "cls_loss"), outputs="total_loss"), PredictBox(width=640, height=640, inputs=("pred_s", "pred_m", "pred_l"), outputs="box_pred", mode="eval"), UpdateOp(model=model, loss_name="total_loss") ]) traces = [ MeanAveragePrecision(num_classes=80, true_key='bbox', pred_key='box_pred', mode="eval"), BestModelSaver(model=model, save_dir=model_dir, metric='mAP', save_best_mode="max") ] lr_schedule = { 1: LRScheduler(model=model, lr_fn=lambda step: lr_schedule_warmup( step, train_steps_epoch=np.ceil(len(train_ds) / batch_size), init_lr=init_lr)), 4: LRScheduler(model=model, lr_fn=lambda epoch: cosine_decay(epoch, cycle_length=epochs - 3, init_lr=init_lr, min_lr=init_lr / 100, start=4)) } traces.append(EpochScheduler(lr_schedule)) estimator = fe.Estimator( pipeline=pipeline, network=network, epochs=epochs, traces=traces, monitor_names=["bbox_loss", "conf_loss", "cls_loss"], train_steps_per_epoch=train_steps_per_epoch, eval_steps_per_epoch=eval_steps_per_epoch) return estimator
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
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(data_dir=None, epochs=12, batch_size_per_gpu=4, im_size=1344, model_dir=tempfile.mkdtemp(), train_steps_per_epoch=None, eval_steps_per_epoch=None): assert im_size % 32 == 0, "im_size must be a multiple of 32" num_device = get_num_devices() train_ds, val_ds = mscoco.load_data(root_dir=data_dir, load_masks=True) batch_size = num_device * batch_size_per_gpu pipeline = fe.Pipeline( train_data=train_ds, eval_data=val_ds, test_data=val_ds, ops=[ ReadImage(inputs="image", outputs="image"), MergeMask(inputs="mask", outputs="mask"), GetImageSize(inputs="image", outputs="imsize", mode="test"), LongestMaxSize(max_size=im_size, image_in="image", mask_in="mask", bbox_in="bbox", bbox_params="coco"), RemoveIf(fn=lambda x: len(x) == 0, inputs="bbox"), PadIfNeeded(min_height=im_size, min_width=im_size, image_in="image", mask_in="mask", bbox_in="bbox", bbox_params="coco", border_mode=cv2.BORDER_CONSTANT, value=0), Sometimes( HorizontalFlip(image_in="image", mask_in="mask", bbox_in="bbox", bbox_params="coco", mode="train")), Resize(height=im_size // 4, width=im_size // 4, image_in='mask'), # downscale mask for memory efficiency Gt2Target(inputs=("mask", "bbox"), outputs=("gt_match", "mask", "classes")), Delete(keys="bbox"), Delete(keys="image_id", mode="!test"), Batch(batch_size=batch_size, pad_value=0) ], num_process=8 * num_device) init_lr = 1e-2 / 16 * batch_size model = fe.build( model_fn=SoloV2, optimizer_fn=lambda x: torch.optim.SGD(x, lr=init_lr, momentum=0.9)) network = fe.Network(ops=[ Normalize(inputs="image", outputs="image", mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), Permute(inputs="image", outputs='image'), ModelOp(model=model, inputs="image", outputs=("feat_seg", "feat_cls_list", "feat_kernel_list")), LambdaOp(fn=lambda x: x, inputs="feat_cls_list", outputs=("cls1", "cls2", "cls3", "cls4", "cls5")), LambdaOp(fn=lambda x: x, inputs="feat_kernel_list", outputs=("k1", "k2", "k3", "k4", "k5")), Solov2Loss(0, 40, inputs=("mask", "classes", "gt_match", "feat_seg", "cls1", "k1"), outputs=("l_c1", "l_s1")), Solov2Loss(1, 36, inputs=("mask", "classes", "gt_match", "feat_seg", "cls2", "k2"), outputs=("l_c2", "l_s2")), Solov2Loss(2, 24, inputs=("mask", "classes", "gt_match", "feat_seg", "cls3", "k3"), outputs=("l_c3", "l_s3")), Solov2Loss(3, 16, inputs=("mask", "classes", "gt_match", "feat_seg", "cls4", "k4"), outputs=("l_c4", "l_s4")), Solov2Loss(4, 12, inputs=("mask", "classes", "gt_match", "feat_seg", "cls5", "k5"), outputs=("l_c5", "l_s5")), CombineLoss(inputs=("l_c1", "l_s1", "l_c2", "l_s2", "l_c3", "l_s3", "l_c4", "l_s4", "l_c5", "l_s5"), outputs=("total_loss", "cls_loss", "seg_loss")), L2Regularizaton(inputs="total_loss", outputs="total_loss_l2", model=model, beta=1e-5, mode="train"), UpdateOp(model=model, loss_name="total_loss_l2"), PointsNMS(inputs="feat_cls_list", outputs="feat_cls_list", mode="test"), Predict(inputs=("feat_seg", "feat_cls_list", "feat_kernel_list"), outputs=("seg_preds", "cate_scores", "cate_labels"), mode="test") ]) train_steps_epoch = int(np.ceil(len(train_ds) / batch_size)) lr_schedule = { 1: LRScheduler( model=model, lr_fn=lambda step: lr_schedule_warmup(step, init_lr=init_lr)), 2: LRScheduler( model=model, lr_fn=lambda step: cosine_decay(step, cycle_length=train_steps_epoch * (epochs - 1), init_lr=init_lr, min_lr=init_lr / 100, start=train_steps_epoch)) } traces = [ EpochScheduler(lr_schedule), COCOMaskmAP(data_dir=val_ds.root_dir, inputs=("seg_preds", "cate_scores", "cate_labels", "image_id", "imsize"), mode="test"), BestModelSaver(model=model, save_dir=model_dir, metric="total_loss") ] estimator = fe.Estimator(pipeline=pipeline, network=network, epochs=epochs, traces=traces, monitor_names=("cls_loss", "seg_loss"), train_steps_per_epoch=train_steps_per_epoch, eval_steps_per_epoch=eval_steps_per_epoch) return estimator
def get_estimator(epochs=20, batch_size=4, train_steps_per_epoch=None, eval_steps_per_epoch=None, save_dir=tempfile.mkdtemp(), log_steps=20, data_dir=None): # step 1 csv = montgomery.load_data(root_dir=data_dir) pipeline = fe.Pipeline( train_data=csv, eval_data=csv.split(0.2), batch_size=batch_size, ops=[ ReadImage(inputs="image", parent_path=csv.parent_path, outputs="image", color_flag='gray'), ReadImage(inputs="mask_left", parent_path=csv.parent_path, outputs="mask_left", color_flag='gray', mode='!infer'), ReadImage(inputs="mask_right", parent_path=csv.parent_path, outputs="mask_right", color_flag='gray', mode='!infer'), CombineLeftRightMask(inputs=("mask_left", "mask_right"), outputs="mask", mode='!infer'), Resize(image_in="image", width=512, height=512), Resize(image_in="mask", width=512, height=512, mode='!infer'), Sometimes(numpy_op=HorizontalFlip( image_in="image", mask_in="mask", mode='train')), Sometimes(numpy_op=Rotate(image_in="image", mask_in="mask", limit=(-10, 10), border_mode=cv2.BORDER_CONSTANT, mode='train')), Minmax(inputs="image", outputs="image"), Minmax(inputs="mask", outputs="mask", mode='!infer') ]) # step 2 model = fe.build( model_fn=lambda: UNet(input_size=(512, 512, 1)), optimizer_fn=lambda: tf.keras.optimizers.Adam(learning_rate=0.0001), model_name="lung_segmentation") network = fe.Network(ops=[ ModelOp(inputs="image", model=model, outputs="pred_segment"), CrossEntropy( inputs=("pred_segment", "mask"), outputs="loss", form="binary"), UpdateOp(model=model, loss_name="loss") ]) # step 3 traces = [ Dice(true_key="mask", pred_key="pred_segment"), BestModelSaver(model=model, save_dir=save_dir, metric='Dice', save_best_mode='max') ] estimator = fe.Estimator(network=network, pipeline=pipeline, epochs=epochs, log_steps=log_steps, 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)