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 instantiate_system(): system = sample_system_object_torch() model = fe.build(model_fn=fe.architecture.pytorch.LeNet, optimizer_fn='adam', model_name='torch') var1 = torch.tensor(1.0) system.network = fe.Network(ops=[ TestTensorOp( inputs="x_out", outputs="x_out", mode="train", var1=var1), ModelOp(model=model, inputs="x_out", outputs="y_pred") ]) system.traces.append(TestTrace(var1=var1)) return system
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 setUpClass(cls): cls.save_dir = tempfile.gettempdir() cls.tf_model = fe.build(model_fn=one_layer_tf_model, optimizer_fn='adam', model_name='tf') cls.torch_model = fe.build(model_fn=MultiLayerTorchModel, optimizer_fn='adam', model_name='torch') cls.data = Data({'loss': 0.5}) cls.state = { 'mode': 'train', 'epoch': 1, 'warmup': False, 'deferred': {}, "scaler": None, "tape": None } cls.tf_input_data = tf.Variable([[2.0, 1.5, 1.0], [1.0, -1.0, -0.5]]) cls.tf_y = tf.constant([[-6], [1]]) cls.torch_input_data = torch.tensor( [[1.0, 1.0, 1.0, -0.5], [0.5, 1.0, -1.0, -0.5]], dtype=torch.float32) cls.torch_y = torch.tensor([[5], [7]], dtype=torch.float32)
def instantiate_system(): system = sample_system_object() model = fe.build(model_fn=fe.architecture.tensorflow.LeNet, optimizer_fn='adam', model_name='tf') var1 = tf.Variable(initial_value=1, trainable=True) system.network = fe.Network(ops=[ TestTensorOp( inputs="x_out", outputs="x_out", mode="train", var1=var1), ModelOp(model=model, inputs="x_out", outputs="y_pred") ]) system.traces.append(TestTrace(var1=var1)) return system
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 get_estimator(batch_size=32, epochs=25, steps_per_epoch=None, validation_steps=None, model_dir=tempfile.mkdtemp()): # load CUB200 dataset. csv_path, path = cub200.load_data() writer = RecordWriter(save_dir=os.path.join(path, "tfrecords"), train_data=csv_path, validation_data=0.2, ops=[ ImageReader(inputs='image', parent_path=path), Resize(target_size=(128, 128), keep_ratio=True, outputs='image'), MatReader(inputs='annotation', parent_path=path), SelectDictKey(), Resize((128, 128), keep_ratio=True), Reshape(shape=(128, 128, 1), outputs="annotation") ]) # data pipeline pipeline = fe.Pipeline(batch_size=batch_size, data=writer, ops=Minmax(inputs='image', outputs='image')) # Network model = fe.build(model_def=UNet, model_name="unet_cub", optimizer=tf.optimizers.Adam(), loss_name="loss") network = fe.Network(ops=[ ModelOp(inputs='image', model=model, outputs='mask_pred'), BinaryCrossentropy( y_true='annotation', y_pred='mask_pred', outputs="loss") ]) # estimator traces = [ Dice(true_key="annotation", pred_key='mask_pred'), ModelSaver(model_name="unet_cub", save_dir=model_dir, save_best=True) ] estimator = fe.Estimator(network=network, pipeline=pipeline, traces=traces, epochs=epochs, steps_per_epoch=steps_per_epoch, validation_steps=validation_steps, log_steps=50) return estimator
def test_estimator_warmup_network_missing_key(self): loader = get_sample_tf_dataset() pipeline = fe.Pipeline(train_data=loader) # "x", "y" model = fe.build(model_fn=LeNetTf, optimizer_fn="adam") network = fe.Network(ops=[ ModelOp(model=model, inputs="x_out", outputs="y_pred") # miss key x_out ]) est = fe.Estimator(pipeline=pipeline, network=network, epochs=1) est._prepare_traces(run_modes={"train", "eval"}) with self.assertRaises(KeyError): est._warmup(warmup=True)
def sample_system_object_torch(): x_train = np.random.rand(3, 28, 28, 3) y_train = np.random.randint(10, size=(3, )) x_eval = np.random.rand(2, 28, 28, 3) y_eval = np.random.randint(10, size=(2, )) train_data = NumpyDataset({'x': x_train, 'y': y_train}) eval_data = NumpyDataset({'x': x_eval, 'y': y_eval}) test_data = eval_data.split(0.5) model = fe.build(model_fn=fe.architecture.pytorch.LeNet, optimizer_fn='adam', model_name='torch') pipeline = fe.Pipeline(train_data=train_data, eval_data=eval_data, test_data=test_data, batch_size=1) network = fe.Network(ops=[ModelOp(model=model, inputs="x_out", outputs="y_pred")]) system = System(network=network, pipeline=pipeline, traces=[], total_epochs=10, mode='train') return system
def fastestimator_run(epochs_pretrain=50, epochs_finetune=10, batch_size=512, train_steps_per_epoch=None, save_dir=tempfile.mkdtemp()): model_con = pretrain_model(epochs_pretrain, batch_size, train_steps_per_epoch, save_dir) model_finetune = fe.build(model_fn=lambda: ResNet9OneLayerHead(length=10), optimizer_fn="adam") model_finetune.encoder.load_state_dict( model_con.encoder.state_dict()) # load the encoder weight finetune_model(model_finetune, epochs_finetune, batch_size, train_steps_per_epoch, save_dir)
def run_test(mixed_precision, merge_grad, gradient): lr = 0.1 lr2 = 0.01 lr3 = 0.001 pipeline = fe.Pipeline(train_data=self.train_data, batch_size=4, ops=[ ExpandDims(inputs="x", outputs="x", axis=0), Minmax(inputs="x", outputs="x") ]) optimizer_fn = EpochScheduler({ 1: lambda x: torch.optim.SGD(params=x, lr=lr), 2: lambda x: torch.optim.SGD(params=x, lr=lr2), 3: lambda x: torch.optim.SGD(params=x, lr=lr3) }) model = fe.build(model_fn=LeNet_torch, optimizer_fn=optimizer_fn, mixed_precision=mixed_precision) network = fe.Network(ops=[ ModelOp(model=model, inputs="x", outputs="y_pred"), CrossEntropy(inputs=("y_pred", "y"), outputs="ce"), GradientOp(model=model, finals="ce", outputs="grad"), UpdateOp(model=model, loss_name="ce", gradients=gradient, merge_grad=merge_grad), ]) traces = [ CheckNetworkWeight(model=model, grad_key="grad", merge_grad=merge_grad, test_self=self, framework="torch", lrs=[lr, lr2, lr3], work_intervals=[[1, 2], [2, 3], [3, 4]]) ] estimator = fe.Estimator(pipeline=pipeline, network=network, epochs=3, traces=traces, train_steps_per_epoch=2) estimator.fit(warmup=False)
def get_estimator(batch_size=256, epochs=50, steps_per_epoch=None, validation_steps=None, model_dir=tempfile.mkdtemp()): # prepare data (x_train, _), (_, _) = tf.keras.datasets.mnist.load_data() data = {"train": {"x": np.expand_dims(x_train, -1)}} pipeline = fe.Pipeline(batch_size=batch_size, data=data, ops=Myrescale(inputs="x", outputs="x")) # prepare model g_femodel = fe.build(model_def=make_generator_model, model_name="gen", loss_name="gloss", optimizer=tf.optimizers.Adam(1e-4)) d_femodel = fe.build(model_def=make_discriminator_model, model_name="disc", loss_name="dloss", optimizer=tf.optimizers.Adam(1e-4)) network = fe.Network(ops=[ ModelOp(inputs=lambda: tf.random.normal([batch_size, 100]), model=g_femodel), ModelOp(model=d_femodel, outputs="pred_fake"), ModelOp(inputs="x", model=d_femodel, outputs="pred_true"), GLoss(inputs=("pred_fake"), outputs="gloss"), DLoss(inputs=("pred_true", "pred_fake"), outputs="dloss") ]) # prepare estimator traces = [ModelSaver(model_name='gen', save_dir=model_dir, save_freq=5)] estimator = fe.Estimator(network=network, pipeline=pipeline, epochs=epochs, traces=traces, steps_per_epoch=steps_per_epoch, validation_steps=validation_steps) return estimator
def get_estimator(epochs=10, batch_size=64, max_len=500, steps_per_epoch=None, validation_steps=None, model_dir=tempfile.mkdtemp()): # step 1. prepare data (x_train, y_train), (x_eval, y_eval) = tf.keras.datasets.imdb.load_data(maxlen=max_len, num_words=MAX_WORDS) data = { "train": { "x": np.array([pad(x, max_len, 0) for x in x_train]), "y": y_train }, "eval": { "x": np.array([pad(x, max_len, 0) for x in x_eval]), "y": y_eval } } pipeline = fe.Pipeline(batch_size=batch_size, data=data, ops=Reshape([1], inputs="y", outputs="y")) # step 2. prepare model model = fe.build(model_def=lambda: create_lstm(max_len), model_name="lstm_imdb", optimizer="adam", loss_name="loss") network = fe.Network(ops=[ ModelOp(inputs="x", model=model, outputs="y_pred"), BinaryCrossentropy(y_true="y", y_pred="y_pred", outputs="loss") ]) traces = [ Accuracy(true_key="y", pred_key="y_pred"), ModelSaver(model_name="lstm_imdb", save_dir=model_dir, save_best=True) ] # step 3.prepare estimator estimator = fe.Estimator(network=network, pipeline=pipeline, epochs=epochs, traces=traces, steps_per_epoch=steps_per_epoch, validation_steps=validation_steps) 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 = cifar10.load_data() pipeline = fe.Pipeline(test_data=train_data, batch_size=batch_size, ops=[Normalize(inputs="x", outputs="x")], num_process=0) weight_path = os.path.abspath(os.path.join(__file__, "..", "resources", "lenet_cifar10_tf.h5")) model = fe.build(model_fn=lambda: LeNet(input_shape=(32, 32, 3)), optimizer_fn="adam", weights_path=weight_path) network = fe.Network(ops=[ ModelOp(model=model, inputs="x", outputs="y_pred") ]) save_dir = tempfile.mkdtemp() traces = [ Saliency(model=model, model_inputs="x", class_key="y", model_outputs="y_pred", samples=5, label_mapping=label_mapping), ImageSaver(inputs="saliency", save_dir=save_dir) ] estimator = fe.Estimator(pipeline=pipeline, network=network, epochs=5, traces=traces, log_steps=1000) estimator.test() ans_img_path = os.path.abspath(os.path.join(__file__, "..", "resources", "saliency_figure.png")) ans_img = img_to_rgb_array(ans_img_path) output_img_path = os.path.join(save_dir, "saliency_test_epoch_5.png") output_img = img_to_rgb_array(output_img_path) self.assertTrue(check_img_similar(output_img, ans_img))
def instantiate_system(): system = sample_system_object_torch() model = fe.build(model_fn=fe.architecture.pytorch.LeNet, optimizer_fn='adam', model_name='torch') system.network = fe.Network(ops=[ RepeatScheduler([ TestTensorOp( inputs="x_out", outputs="x_out", mode="train", var=1), TestTensorOp( inputs="x_out", outputs="x_out", mode="train", var=1) ]), ModelOp(model=model, inputs="x_out", outputs="y_pred") ]) return system
def build_model(self): """ Define the FastEstimator model architecture. Args: None Returns: model: Union[tf.keras.sequential, nn.module] """ model = fe.build(model_fn=lambda: LeNet(input_shape=(32, 32, 3)), optimizer_fn="adam", model_name="adv_model") return model
def test_tf_model_end_to_end_gradient(self): train_data, _ = mnist.load_data() pipeline = fe.Pipeline(train_data=train_data, batch_size=4, ops=[ExpandDims(inputs="x", outputs="x"), Minmax(inputs="x", outputs="x")]) model = fe.build(model_fn=LeNet_tf, optimizer_fn="adam") network = fe.Network(ops=[ ModelOp(model=model, inputs="x", outputs="y_pred"), CrossEntropy(inputs=("y_pred", "y"), outputs="ce"), GradientOp(model=model, finals="ce", outputs="gradients"), UpdateOp(model=model, gradients="gradients", loss_name="ce") ]) estimator = fe.Estimator(pipeline=pipeline, network=network, epochs=2, max_train_steps_per_epoch=10) estimator.fit()
def instantiate_system(): system = sample_system_object() model = fe.build(model_fn=fe.architecture.tensorflow.LeNet, optimizer_fn='adam', model_name='tf') system.network = fe.Network(ops=[ fe.op.tensorop.meta.Fuse(ops=[ TestTensorOp( inputs="x_out", outputs="x_out", mode="train", var=1), TestTensorOp( inputs="x_out", outputs="x_out", mode="train", var=1) ]), ModelOp(model=model, inputs="x_out", outputs="y_pred") ]) return system
def test_estimator_warmup_torch_dataset_tf_model_smoke(self): loader = get_sample_torch_dataloader(expand_axis=-1) pipeline = fe.Pipeline(train_data=loader) # "x", "y" model = fe.build(model_fn=LeNetTf, optimizer_fn="adam") network = fe.Network( ops=[ModelOp(model=model, inputs="x", outputs="y_pred")]) est = fe.Estimator(pipeline=pipeline, network=network, epochs=1, traces=[Trace(inputs="y_pred")]) est._prepare_traces(run_modes={"train", "eval"}) est._warmup(warmup=True) self.assertTrue(True)
def test_torch_model_with_arbitrary_gradient(self): lr = 0.1 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model = fe.build(model_fn=OneLayerTorchModel, optimizer_fn=lambda x: torch.optim.SGD(params=x, lr=lr)).to(device) init_weights = [deepcopy(x).cpu().detach().numpy() for x in model.parameters() if x.requires_grad] gradients = [torch.tensor([[1.0, 1.0, 1.0]]).to(torch.device(device))] fe.backend.update_model(model, gradients=gradients) gradients = [x.cpu().detach().numpy() for x in gradients] new_weights = [x.cpu().detach().numpy() for x in model.parameters() if x.requires_grad] for init_w, new_w, grad in zip(init_weights, new_weights, gradients): new_w_ans = init_w - grad * lr self.assertTrue(np.allclose(new_w_ans, new_w))
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(data_dir=None, model_dir=tempfile.mkdtemp(), epochs=20, em_dim=128, batch_size=64, train_steps_per_epoch=None, 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, ops=[ Encode(inputs="source", outputs="source", tokenizer=pt_tokenizer), Encode(inputs="target", outputs="target", tokenizer=en_tokenizer), Batch(batch_size=batch_size, pad_value=0) ]) model = fe.build( model_fn=lambda: transformer(num_layers=4, em_dim=em_dim, num_heads=8, dff=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_mask", "decode_mask")), ModelOp(model=model, inputs=("source", "target_inp", "encode_mask", "decode_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, train_steps_per_epoch=train_steps_per_epoch, eval_steps_per_epoch=eval_steps_per_epoch) 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 test_tf_on_batch_end(self): tensorboard = TensorBoard(log_dir=self.log_dir, weight_histogram_freq=1, update_freq=1) tensorboard.system = sample_system_object() tensorboard.system.global_step = 1 tensorboard.writer = _TfWriter(self.log_dir, '', tensorboard.system.network) model = fe.build(model_fn=fe.architecture.tensorflow.LeNet, optimizer_fn='adam') tensorboard.system.network.epoch_models = {model} if os.path.exists(self.train_path): shutil.rmtree(self.train_path) tensorboard.on_batch_end(data=self.tf_data) filepath = getfilepath() for e in tf.compat.v1.train.summary_iterator(filepath): for v in e.summary.value: if v.tag == "tf_dense_1/bias_0": output = v.histo.num self.assertEqual(output, 10.0)
def test_estimator_warmup_trace_missing_key(self): loader = get_sample_tf_dataset() pipeline = fe.Pipeline(train_data=loader) # "x", "y" model = fe.build(model_fn=LeNetTf, optimizer_fn="adam") network = fe.Network(ops=[ EpochScheduler({ 1: ModelOp(model=model, inputs="x", outputs="y_pred"), 2: ModelOp(model=model, inputs="x", outputs="y_pred") }) # miss key x_out ]) est = fe.Estimator(pipeline=pipeline, network=network, epochs=2, traces=[Trace(inputs="z")]) # miss key "z" est._prepare_traces(run_modes={"train", "eval"}) with self.assertRaises(AssertionError): est._warmup()
def test_torch_model_on_epoch_end_reduce_lr(self): model = fe.build(model_fn=MultiLayerTorchModel, optimizer_fn='adam') model_name = model.model_name + '_lr' lr_on_plateau = ReduceLROnPlateau(model=model, metric='loss') lr_on_plateau.system = sample_system_object() lr_on_plateau.best = 5 lr_on_plateau.wait = 11 with patch('sys.stdout', new=StringIO()) as fake_stdout: lr_on_plateau.on_epoch_end(data=self.data) log = fake_stdout.getvalue().strip() self.assertEqual(log, self.torch_expected_msg) with self.subTest('Check learning rate in data'): self.assertTrue( math.isclose(self.data[model_name], 0.000100000005, rel_tol=1e-3))
def get_estimator(): pos_real, _ = mnist.load_data() neg_real, _ = mnist.load_data() neg_sim, _ = mnist.load_data() neg_sim = NegativeImageSimulatedTube(neg_sim) batch_ds = BatchDataset(datasets=(pos_real, neg_real, neg_sim), num_samples=(2, 2, 1)) pipeline = fe.Pipeline(train_data=batch_ds, ops=[ExpandDims(inputs="x", outputs="x"), Minmax(inputs="x", outputs="x")]) model = fe.build(model_fn=LeNet, optimizer_fn="adam") network = fe.Network(ops=[ ModelOp(model=model, inputs="x", outputs="y_pred"), CrossEntropy(inputs=("y_pred", "y"), outputs="ce"), UpdateOp(model=model, loss_name="ce") ]) estimator = fe.Estimator(pipeline=pipeline, network=network, epochs=2) return estimator
def test_tf_model_with_loss(self): def update(x, model): with tf.GradientTape(persistent=True) as tape: y = fe.backend.feed_forward(model, x) fe.backend.update_model(model, loss=y, tape=tape) model = fe.build(model_fn=one_layer_tf_model, optimizer_fn="adam") init_weight = get_tf_model_weight(model) x = tf.constant([[1, 1, 1]]) strategy = tf.distribute.get_strategy() if isinstance(strategy, tf.distribute.MirroredStrategy): strategy.run(update, args=(x, model)) else: update(x, model) new_weight = get_tf_model_weight(model) self.assertTrue(not is_equal(init_weight, new_weight))
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))