train_har.py

# The testing module requires faiss
from functools import partial
import sys
# So if you don't have that, then this import will break
from pml import trainers
from pml import losses, miners, samplers, testers, utils
# import losss
import torch.nn as nn
import record_keeper
from utils import common_functions as c_f
import pml.utils.logging_presets as logging_presets
import pml
import pml as pytorch_metric_learning
from torchvision import datasets, models, transforms
import torchvision
import logging

logging.getLogger().setLevel(logging.INFO)
import os
# from pytorch_pretrained_vit import ViT
from pml.losses.base_metric_loss_function import BaseMetricLossFunction
from pml.testers.base_tester import BaseTester
from vit_pytorch.pvt import PyramidVisionTransformer
from vit_pytorch.CausalLevit import LeViT_384

logging.info("pytorch-metric-learning VERSION %s" % pytorch_metric_learning.__version__)
logging.info("record_keeper VERSION %s" % record_keeper.__version__)

from sklearn.metrics import accuracy_score
# from efficientnet_pytorch import EfficientNet
import torch
import numpy as np
import pickle

import hydra
from omegaconf import DictConfig

# reprodcibile
np.random.seed(42)
torch.manual_seed(42)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False


class Loss_n(BaseMetricLossFunction):
    def __init__(self):
        super().__init__()
        self.n_pair_loss = losses.NPairsLoss()
        self.angular_loss = losses.AngularLoss(alpha=45)

    def compute_loss(self, embeddings, labels, indices_tuple):
        dict_angular = self.angular_loss.compute_loss(embeddings, labels, indices_tuple)
        dict_npair = self.n_pair_loss.compute_loss(embeddings, labels, indices_tuple)
        losses = 0.01 * dict_angular['loss']['losses'] + 0.1 * dict_npair['loss']['losses']
        dict_angular['loss']['losses'] = losses
        return dict_angular


def calibration_augmentation(base_means, base_cov, embedding_and_labels):
    n_shot = 1
    n_ways = 20
    support_data = np.nan_to_num(embedding_and_labels['samples'][0])
    support_label = embedding_and_labels['samples'][1]
    sampled_data = []
    sampled_label = []
    num_sampled = int(10 / n_shot)
    for i in range(20):
        mean, cov = distribution_calibration(support_data[np.squeeze(support_label == i), :], base_means, base_cov, k=4)
        sampled_data.append(np.random.multivariate_normal(mean=mean, cov=cov, size=num_sampled))
        sampled_label.extend([support_label[i]] * num_sampled)
    sampled_data = np.concatenate([sampled_data[:]]).reshape(n_ways * n_shot * num_sampled, -1)
    X_aug = np.concatenate([support_data, sampled_data])
    Y_aug = np.concatenate([support_label, sampled_label])
    return X_aug, Y_aug


def distribution_calibration(query, base_means, base_cov, k, alpha=0.21):
    dist = []
    for i in range(len(base_means)):
        dist.append(np.linalg.norm(query - base_means[i]))
    index = np.argpartition(dist, k)[:k]

    mean = np.concatenate([np.array(base_means)[index], query])
    calibrated_mean = np.mean(mean, axis=0)
    # print(base_cov)
    calibrated_cov = np.mean(np.array(base_cov)[index], axis=0) + alpha

    return calibrated_mean, calibrated_cov


class OneShotTester(BaseTester):

    def __init__(self, end_of_testing_hook=None):
        super().__init__()
        self.max_accuracy = 0.0
        self.embedding_filename = ""
        self.end_of_testing_hook = end_of_testing_hook

    def __get_correct(self, output, target, topk=(1,)):
        with torch.no_grad():
            maxk = max(topk)
            batch_size = target.size(0)

            _, pred = output.topk(maxk, 1, True, True)
            pred = pred.t()
            correct = pred.eq(target.view(1, -1).expand_as(pred))
            #             print(correct)
        return correct

    def __accuracy(self, output, target, topk=(1,)):
        """Computes the accuracy over the k top predictions for the specified values of k"""
        with torch.no_grad():
            correct = self.__get_correct(output, target, topk)
            batch_size = target.size(0)
            res = []
            for k in topk:
                correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
                res.append(correct_k.mul_(100.0 / batch_size))
            return res

    def do_knn_and_accuracies(self, accuracies, embeddings_and_labels, split_name, classifier_model, tag_suffix='',):
        # print(embeddings_and_labels)
        # train_embeddings = embeddings_and_labels['train'][0]
        # train_labels = embeddings_and_labels['train'][1]
        # print(train_embeddings.shape)
        # print(train_labels.shape)
        query_prediction = torch.argmax(torch.nn.functional.softmax(classifier_model(torch.from_numpy(embeddings_and_labels["val"][0]).cuda()), dim=-1), dim=-1).cpu().numpy()
        query_labels = embeddings_and_labels["val"][1]
        #reference_embeddings = embeddings_and_labels["samples"][0]
        #reference_labels = embeddings_and_labels["samples"][1]
        #knn_indices, knn_distances = utils.stat_utils.get_knn(reference_embeddings, query_embeddings, 1, False)
        #knn_labels = reference_labels[knn_indices][:, 0]

        accuracy = accuracy_score(query_prediction, query_labels)
        print(accuracy)
        #with open(self.embedding_filename + "_last_ms_levit", 'wb') as f:
        #    print("Dumping embeddings for new max_acc to file", self.embedding_filename + "_last")
        #    pickle.dump([query_embeddings, query_labels, reference_embeddings, reference_labels, accuracy], f)
        accuracies["accuracy"] = accuracy
        keyname = self.accuracies_keyname("mean_average_precision_at_r")  # accuracy as keyname not working
        accuracies[keyname] = accuracy

    def do_knn_and_accuracies_aug(self, accuracies, embeddings_and_labels, split_name, tag_suffix=''):
        # print(embeddings_and_labels)
        print("test")
        train_embeddings = embeddings_and_labels['train'][0]
        train_labels = embeddings_and_labels['train'][1]
        # print(train_labels.shape)
        # print(train_embeddings.shape)
        base_means = []
        base_cov = []
        for key in range(100):
            feature = train_embeddings[np.squeeze(train_labels == key, axis=1), :]
            mean = np.mean(feature, axis=0)
            cov = np.cov(feature.T)
            base_means.append(mean)
            base_cov.append(cov)
            # print(cov.shape)
        x_aug, y_aug = calibration_augmentation(base_means, base_cov, embeddings_and_labels)
        # print(train_embeddings.shape)
        # print(train_labels.shape)
        reference_embedding = []
        reference_labels = []
        for key in range(20):
            reference_embedding.append(x_aug[np.squeeze(y_aug == key), :].mean(0))
            reference_labels.append(key)
        reference_embeddings = np.stack(reference_embedding, axis=0)
        reference_labels = np.stack(reference_labels, axis=0)
        query_embeddings = embeddings_and_labels["val"][0]
        query_labels = embeddings_and_labels["val"][1]
        # reference_embeddings = #embeddings_and_labels["samples"][0]
        # reference_labels = #embeddings_and_labels["samples"][1]
        # print(reference_embeddings.shape)
        # print(query_embeddings.shape)
        knn_indices, knn_distances = utils.stat_utils.get_knn(reference_embeddings.astype(np.float32), query_embeddings,
                                                              1, False)
        knn_labels = reference_labels[knn_indices][:, 0]
        accuracy = accuracy_score(knn_labels, query_labels)
        print(accuracy)
        accuracies["accuracy"] = accuracy
        keyname = self.accuracies_keyname("mean_average_precision_at_r")  # accuracy as keyname not working
        accuracies[keyname] = accuracy


class MLP(nn.Module):
    # layer_sizes[0] is the dimension of the input
    # layer_sizes[-1] is the dimension of the output
    def __init__(self, layer_sizes, final_relu=False):
        super().__init__()
        layer_list = []
        layer_sizes = [int(x) for x in layer_sizes]
        num_layers = len(layer_sizes) - 1
        final_relu_layer = num_layers if final_relu else num_layers - 1
        for i in range(len(layer_sizes) - 1):
            input_size = layer_sizes[i]
            curr_size = layer_sizes[i + 1]
            if i < final_relu_layer:
                layer_list.append(nn.ReLU(inplace=True))
            layer_list.append(nn.Linear(input_size, curr_size))
        self.net = nn.Sequential(*layer_list)
        self.last_linear = self.net[-1]

    def forward(self, x):
        out = self.net(x)
        # print(out.size())
        return out


class SignalModule(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        # self.trunk_signal = torchvision.models.__dict__[cfg.model.model_name](pretrained=cfg.model.pretrained)
        # ViT(image_size=256,patch_size=64, num_classes=21, dim=512,depth=6, heads=16, mlp_dim=21,dropout=0.1,emb_dropout=0.1)

        self.trunk_signal = LeViT_384(num_classes=512, distillation=True,
                                      pretrained=False, fuse=False)
        self.trunk_fft = LeViT_384(num_classes=512, distillation=True,
                                   pretrained=False, fuse=False, in_chans=6)

        """PyramidVisionTransformer(
            patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], qkv_bias=True,
            norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1],
        )
        self.trunk_fft = PyramidVisionTransformer(
            patch_size=4, in_chans=6, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], qkv_bias=True,
            norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1],
        )"""

        """
        self.trunk_fft = LeViT(
            image_size = 256,
            num_classes = 20,
            stages = 3,             # number of stages
            dim = (256, 384, 512),  # dimensions at each stage
            depth = 4,              # transformer of depth 4 at each stage
            heads = (4, 6, 8),      # heads at each stage
            mlp_mult = 2,
            dropout = 0.1
         )
        """

        # self.trunk_signal.fc = Identity()
        # self.trunk_fft.fc =Identity()
        # ViT(image_size=256,patch_size=64, num_classes=21, dim=512,depth=6, heads=16,channels=6, mlp_dim=21,dropout=0.1,emb_dropout=0.1)
        # self.conv_fusion = nn.Sequential(
        #    nn.Conv2d(1024, 512, kernel_size=3, padding=1, bias=False),
        #    nn.BatchNorm2d(1024, eps=1e-3, momentum=0.01),
        #    nn.ReLU()
        # )
        self.MLP_2 = MLP([1024, 512])
        # self.MLP = MLP([6,3])

    def forward(self, x):
        # print(x.size())
        batch_size = x.size()[0]
        c_1 = self.trunk_signal(x)

        data_fft = torch.rfft(x.permute(0, 1, 3, 2), signal_ndim=1)
        # print(data_fft[:,:,:,:,1])
        # sys.exit()
        real = data_fft[:, :, :, :, 0].permute(0, 1, 3, 2)
        imag = data_fft[:, :, :, :, 1].permute(0, 1, 3, 2)
        norm = torch.nn.functional.normalize(torch.sqrt(torch.pow(real, 2) + torch.pow(imag, 2)), dim=-2)
        angle = torch.nn.functional.normalize(torch.atan2(real, imag), dim=-2)
        data_fusion = torch.cat([norm, angle], dim=1)
        # print(data_fusion.size())
        # sys.exit()
        container = torch.zeros([batch_size, 6, 256, 256]).cuda()
        container[:, :, :129, :] = data_fusion[:, :, :, :]
        container[:, :, 129:, :] = torch.flip(container, dims=[2])[:, :, :127, :]
        c_2 = self.trunk_fft(container)
        c_2 = self.trunk_signal(x)
        # print(c_1.size(), c2.size())
        fusion = self.MLP_2(torch.cat([c_1, c_2], dim=1))

        return fusion


# This is for replacing the last layer of a pretrained network.
# This code is from https://github.com/KevinMusgrave/powerful_benchmarker
class Identity(nn.Module):

    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x


def get_datasets(data_dir, cfg, mode="train"):
    common_transforms = []
    train_transforms = []
    test_transforms = []
    # if cfg.transform.transform_resize_match:
    common_transforms.append(transforms.Resize((256, 256)))

    if cfg.transform.transform_random_resized_crop:
        train_transforms.append(transforms.RandomResizedCrop(cfg.transform.transform_resize))
    if cfg.transform.transform_random_horizontal_flip:
        train_transforms.append(torchvision.transforms.RandomHorizontalFlip(p=0.5))
    if cfg.transform.transform_random_rotation:
        train_transforms.append(
            transforms.RandomRotation(cfg.transform.transform_random_rotation_degrees))  # , fill=255))
    if cfg.transform.transform_random_shear:
        train_transforms.append(torchvision.transforms.RandomAffine(0,
                                                                    shear=(
                                                                        cfg.transform.transform_random_shear_x1,
                                                                        cfg.transform.transform_random_shear_x2,
                                                                        cfg.transform.transform_random_shear_y1,
                                                                        cfg.transform.transform_random_shear_y2
                                                                    ),
                                                                    fillcolor=255))
    if cfg.transform.transform_random_perspective:
        train_transforms.append(transforms.RandomPerspective(distortion_scale=cfg.transform.transform_perspective_scale,
                                                             p=0.5,
                                                             interpolation=3)
                                )
    if cfg.transform.transform_random_affine:
        train_transforms.append(transforms.RandomAffine(degrees=(cfg.transform.transform_degrees_min,
                                                                 cfg.transform.transform_degrees_max),
                                                        translate=(cfg.transform.transform_translate_a,
                                                                   cfg.transform.transform_translate_b),
                                                        fillcolor=255))
    data_transforms = {
        'train': transforms.Compose(common_transforms + train_transforms + [transforms.ToTensor()]),
        'test': transforms.Compose(common_transforms + [transforms.ToTensor()]),
    }

    train_dataset = datasets.ImageFolder(os.path.join(data_dir, "train"),
                                         data_transforms["train"])
    print(train_dataset.class_to_idx)

    # for the final model we can join train, validation, validation samples datasets
    print(mode)
    if mode == "final_train":
        # train_dataset = torch.utils.data.ConcatDataset([train_dataset,
        #        val_dataset,
        #        val_samples_dataset])

        test_dataset = datasets.ImageFolder(os.path.join(data_dir, "test"),
                                            data_transforms["test"])
        print(test_dataset.class_to_idx)


        return train_dataset, test_dataset
    else:
        if mode == "train":
            val_dataset = datasets.ImageFolder(os.path.join(data_dir, "val"),
                                               data_transforms["test"])


            return train_dataset, val_dataset

        if mode == "test":
            return train_dataset, test_dataset


@hydra.main(config_path="config/config.yaml")
def train_app(cfg):
    print(cfg.pretty())

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # trunk = PyramidVisionTransformer(
    #    patch_size=4, embed_dims=[128, 256, 512, 768], num_heads=[2, 4, 8, 12], mlp_ratios=[8, 8, 4, 4], qkv_bias=True,
    #    norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 10, 60, 3], sr_ratios=[8, 4, 2, 1],
    #     )
    # trunk = PyramidVisionTransformer(
    #    patch_size=32, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], qkv_bias=True,
    #    norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1],
    #    )
    trunk = LeViT_384(num_classes=512, distillation=True,
                      pretrained=False, fuse=False)
    # trunk = LeViT_384_cascaded(num_classes=512, distillation=True,
    #          pretrained=False, fuse=False)
    # device=torch.device("cpu")
    # Set trunk model and replace the softmax layer with an identity function
    # trunk = torchvision.models.__dict__[cfg.model.model_name](pretrained=cfg.model.pretrained)
    # trunk = SignalModule(cfg)
    # trunk=ViT(image_size=256,patch_size=64, num_classes=21, dim=512,depth=6, heads=16, mlp_dim=21,dropout=0.1,e

    # trunk= torchvision.models.__dict__[cfg.model.model_name](pretrained=cfg.model.pretrained)
    # trunk = models.resnet18(pretrained=False)
    # trunk = models.alexnet(pretrained=True)
    # trunk = models.resnet50(pretrained=True)
    # trunk = models.resnet152(pretrained=True)
    # trunk = models.wide_resnet50_2(pretrained=True)
    # trunk = EfficientNet.from_pretrained('efficientnet-b2')
    # trunk = ViT('B_16_imagenet1k', pretrained=True)
    trunk_output_size = 512
    trunk.fc = Identity()
    trunk = torch.nn.DataParallel(trunk.to(device))

    embedder = torch.nn.DataParallel(MLP([trunk_output_size, cfg.embedder.size]).to(device))
    classifier = torch.nn.DataParallel(MLP([cfg.embedder.size, 120])).to(device)

    # Set optimizers
    if cfg.optimizer.name == "sdg":
        trunk_optimizer = torch.optim.SGD(trunk.parameters(), lr=cfg.optimizer.lr, momentum=cfg.optimizer.momentum,
                                          weight_decay=cfg.optimizer.weight_decay)
        embedder_optimizer = torch.optim.SGD(embedder.parameters(), lr=cfg.optimizer.lr,
                                             momentum=cfg.optimizer.momentum, weight_decay=cfg.optimizer.weight_decay)
        classifier_optimizer = torch.optim.SGD(classifier.parameters(), lr=cfg.optimizer.lr,
                                               momentum=cfg.optimizer.momentum, weight_decay=cfg.optimizer.weight_decay)
    elif cfg.optimizer.name == "rmsprop":
        trunk_optimizer = torch.optim.RMSprop(trunk.parameters(), lr=cfg.optimizer.lr, momentum=cfg.optimizer.momentum,
                                              weight_decay=cfg.optimizer.weight_decay)
        embedder_optimizer = torch.optim.RMSprop(embedder.parameters(), lr=cfg.optimizer.lr,
                                                 momentum=cfg.optimizer.momentum,
                                                 weight_decay=cfg.optimizer.weight_decay)
        classifier_optimizer = torch.optim.RMSprop(classifier.parameters(), lr=cfg.optimizer.lr,
                                                   momentum=cfg.optimizer.momentum,
                                                   weight_decay=cfg.optimizer.weight_decay)
    elif cfg.optimizer.name == 'adam':
        trunk_optimizer = torch.optim.Adam(trunk.parameters(), lr=cfg.optimizer.lr,
                                           weight_decay=cfg.optimizer.weight_decay)
        embedder_optimizer = torch.optim.Adam(embedder.parameters(), lr=cfg.optimizer.lr,
                                              weight_decay=cfg.optimizer.weight_decay)
        classifier_optimizer = torch.optim.Adam(classifier.parameters(), lr=cfg.optimizer.lr,
                                                weight_decay=cfg.optimizer.weight_decay)

    # Set the datasets
    data_dir = os.environ["DATASET_FOLDER"]
    print("Data dir: " + data_dir)

    train_dataset, val_dataset = get_datasets(data_dir, cfg, mode=cfg.mode.type)
    print("Trainset: ", len(train_dataset), "Testset: ", len(val_dataset))

    # Set the loss function
    if cfg.embedder_loss.name == "margin_loss":
        loss = losses.MarginLoss(margin=cfg.embedder_loss.margin, nu=cfg.embedder_loss.nu, beta=cfg.embedder_loss.beta)
    # if cfg.embedder_loss.name == "triplet_margin":
    loss = losses.TripletMarginLoss(margin=cfg.embedder_loss.margin)
    # loss_angular = losses.AngularLoss(alpha=40)
    if cfg.embedder_loss.name == "multi_similarity":
        loss = losses.MultiSimilarityLoss(alpha=cfg.embedder_loss.alpha, beta=cfg.embedder_loss.beta,
                                          base=cfg.embedder_loss.base)
    # if cfg.embedder_loss.name == "proxyanchor":
    # loss = Loss() #losses.ProxyAnchorLoss(num_classes = 22, embedding_size = cfg.embedder.size).cuda()
    # Set the classification loss:
    classification_loss = torch.nn.CrossEntropyLoss()

    # Set the mining function

    if cfg.miner.name == "triplet_margin":
        # miner = miners.TripletMarginMiner(margin=0.2)
        miner = miners.TripletMarginMiner(margin=cfg.miner.margin)
    if cfg.miner.name == "multi_similarity":
        miner = miners.MultiSimilarityMiner(epsilon=cfg.miner.epsilon)
        # miner = miners.MultiSimilarityMiner(epsilon=0.05)

    batch_size = cfg.trainer.batch_size
    num_epochs = cfg.trainer.num_epochs
    iterations_per_epoch = cfg.trainer.iterations_per_epoch
    # Set the dataloader sampler
    sampler = samplers.MPerClassSampler(train_dataset.targets, m=4, length_before_new_iter=len(train_dataset))

    # Package the above stuff into dictionaries.
    models = {"trunk": trunk, "embedder": embedder, "classifier": classifier}
    optimizers = {"trunk_optimizer": trunk_optimizer, "embedder_optimizer": embedder_optimizer,
                  "classifier_optimizer": classifier_optimizer}
    loss_funcs = {"metric_loss": loss, "classifier_loss": classification_loss}
    mining_funcs = {"tuple_miner": miner}

    # We can specify loss weights if we want to. This is optional
    loss_weights = {"metric_loss": cfg.loss.metric_loss, "classifier_loss": cfg.loss.classifier_loss}

    schedulers = {
        # "metric_loss_scheduler_by_epoch": torch.optim.lr_scheduler.StepLR(classifier_optimizer, cfg.scheduler.step_size, gamma=cfg.scheduler.gamma),
        "embedder_scheduler_by_epoch": torch.optim.lr_scheduler.StepLR(embedder_optimizer, 10,
                                                                       gamma=cfg.scheduler.gamma),
        "classifier_scheduler_by_epoch": torch.optim.lr_scheduler.StepLR(classifier_optimizer, 10,
                                                                         gamma=cfg.scheduler.gamma),
        "trunk_scheduler_by_epoch": torch.optim.lr_scheduler.StepLR(embedder_optimizer, 10, gamma=cfg.scheduler.gamma),
    }  # cfg.scheduler.step_size

    experiment_name = "model_c2_cat_levit_%s_model_%s_cl_%s_ml_%s_miner_%s_mix_ml_%02.2f_mix_cl_%02.2f_resize_%d_emb_size_%d_class_size_%d_opt_%s_lr_%02.2f_m_%02.2f_wd_%02.2f" % (
    cfg.dataset.name,
    cfg.model.model_name,
    "cross_entropy",
    cfg.embedder_loss.name,
    cfg.miner.name,
    cfg.loss.metric_loss,
    cfg.loss.classifier_loss,
    cfg.transform.transform_resize,
    cfg.embedder.size,
    cfg.embedder.class_out_size,
    cfg.optimizer.name,
    cfg.optimizer.lr,
    cfg.optimizer.momentum,
    cfg.optimizer.weight_decay)
    record_keeper, _, _ = logging_presets.get_record_keeper("logs_c_2_cat/%s" % (experiment_name),
                                                            "tensorboard_c_2_cat/%s" % (experiment_name))
    hooks = logging_presets.get_hook_container(record_keeper)
    dataset_dict = {"val": val_dataset}
    model_folder = "example_saved_models_c_2_cat/%s/" % (experiment_name)

    # Create the tester
    tester = OneShotTester(
        end_of_testing_hook=hooks.end_of_testing_hook,
        # size_of_tsne=20
    )
    # tester.embedding_filename=data_dir+"/embeddings_pretrained_triplet_loss_multi_similarity_miner.pkl"
    tester.embedding_filename = data_dir + "/" + experiment_name + ".pkl"
    end_of_epoch_hook = hooks.end_of_epoch_hook(tester, dataset_dict, model_folder)
    trainer = trainers.TrainWithClassifier(models,
                                           optimizers,
                                           batch_size,
                                           loss_funcs,
                                           mining_funcs,
                                           train_dataset,
                                           sampler=sampler,
                                           lr_schedulers=schedulers,
                                           dataloader_num_workers=cfg.trainer.batch_size,
                                           loss_weights=loss_weights,
                                           end_of_iteration_hook=hooks.end_of_iteration_hook,
                                           end_of_epoch_hook=end_of_epoch_hook
                                           )

    trainer.train(num_epochs=num_epochs)

    tester = OneShotTester()


if __name__ == "__main__":
    train_app()