# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Training MoCoGAN
#
# In this notebook, the process of using the class Trainer will be explained and the training of MoCoGAN will be carried on as if the `newTrain.py` script would be launched from Console.

# ## Import Models and Dataloading
# First, let's import classes from `models` and `dataloading` modules.
#
# Then let's create the dataloader and define the clParameters.

import sys
sys.path.append('./mocogan')

# +
from mocogan.dataloading import DataLoaderFactory, Transformator
from mocogan.models      import UCF_101, Discriminator_I, VideoDiscriminator, VideoGenerator

def createModels(parameters, dict_For_Classes = None):

    n_classes  = len(parameters['classes'])

    n_channels = 3
    dim_z_content=50                                            #  dimensionality of the content input, ie hidden space [default: 50]
    dim_z_motion=10                                             #  dimensionality of the motion input [default: 10]
    dim_z_category= n_classes if n_classes > 0 else 101         #  dimensionality of categorical input [default: 101]
    video_lenght = 16                                           #  hyperparameter T of the mocogan [default: 16]
    cuda = parameters['cuda']

    return Discriminator_I(), VideoDiscriminator(n_channels, dim_z_category, cuda = cuda), \
                 VideoGenerator(n_channels, dim_z_content, dim_z_category, dim_z_motion, video_lenght, cuda, class_to_idx = dict_For_Classes)


img_size          = 96
max_frame         = 16
mean              = (100.99800554447337/255, 96.7195209000943/255, 89.63882431650443/255)
std               = (72.07041943699456/255, 70.41506399740703/255, 71.55581999303428/255)

batch_size        = 16
classes           = ['Biking', 'HorseRace', 'HorseRiding', 'Bowling', 'IceDancing']
preprocessing     = 0
cuda              = 1
soft_labels       = True
random_labels     = True
shuffle_labels    = True
i_wasserstein     = 5
n_iter            = 100
i_log_stat        = 1
pre_train_epoch   = 0
i_save_weights    = i_wasserstein
i_alternate_train = 2



parameters = {
        "cuda"              : cuda,
        "batch_size"        : batch_size,
        "soft_labels"       : soft_labels,
        "n_iter"            : n_iter,
        "classes"           : classes,
        "i_log_stat"        : i_log_stat,
        "pre_train_epoch"   : pre_train_epoch,
        "i_save_weights"    : i_save_weights,
        "i_alternate_train" : i_alternate_train,
        "preprocessed"      : preprocessing,
        "random_labels"     : random_labels,
        "shuffle_labels"    : shuffle_labels,
        "i_wasserstein"     : i_wasserstein
    }





transformator = Transformator(img_size, max_frame, medium= mean, stdDev= std)
dataloader    = DataLoaderFactory(UCF_101, transformator, preprocessing, batch_size, classes = classes).getDataLoader()

dis_i, dis_v, gen = createModels(parameters, dataloader.dataset.class_to_idx)

gen.init_weigths()
# -

# ## Create the Trainer
# Let's import and use the Trainer Object to train the networks.

# +
from mocogan.trainer import Trainer

trainer = Trainer(parameters, dis_i, dis_v, gen, dataloader, transformator)
# -

# ## Start Training
# Finally, let's start training process in the `Trainer`

trainer.train()

StarcoderdataJupyterScriptsDedupFiltered