Skip to content

Configuring the config

Maksim Nikolaev edited this page Aug 28, 2021 · 1 revision

On this page, I will show you how to configure the config for training and testing any models of the added models.

For convenience, working presets were created, which are located in the configs folder, for example, let's analyze the StyleGAN config:

{
    "Generator": "StyleGAN",
    "Discriminator": "StyleGAN",
    "Loss_gen": "r1",
    "Loss_disc": "r1",
    "Optim_G": "RMSprop",
    "Optim_D": "RMSprop",
    "Gen_config": {
        "res": 256,
        "latent_size": 256,
        "deep_mapping": 8
    },
    "Disc_config": {
        "res": 256
    },
    "Optim_G_config": {
        "lr": 0.0001,
        "alpha": 0.99
    },
    "Optim_D_config": {
        "lr": 0.0001,
        "alpha": 0.99
    },
    "Loss_config": {
        "lambda_gp": 10
    },
    "Trainer": "gp",
    "IMG_SIZE": 256,
    "z_dim": 256,
    "BATCH_SIZE": 64,
    "UPD_FOR_GEN": 1,
    "epochs": 10000,
    "Dataset": "data/DataGAN256",
    "Weight_dir": "Weight"
}

Let's analyze all the required parameters that should be in all configs:

  • Generator - the name of the generator from the added models, for information on how to add new models, see Universal Loader wiki page.

  • Discriminator - the name of the discriminator from the added models.

  • Loss_gen/Loss_disc - the name of the loss functions for the generator and discriminator, respectively.

  • Optim_G/Optim_D - the name for the generator and discriminator optimizers, respectively, which are loaded from torch.optim.

  • Gen_config/Disc_config - these are dictionaries that are passed as kwargs to the generator and discriminator classes, respectively, for their initialization.

  • Optim_G_config/Optim_D_config - similarly, it passes parameters in the form of kwargs for initializing optimizers.

  • Trainer - the corresponding name of one of the classes in trainer.py, for specific loss functions or generators. For the simplest models, the universal "base" is used, for loss functions such as WGAN-GP, "gp" is used.

  • IMG_SIZE - a separate parameter that means the size of the generated image for the case if it is passed in the model parameters.

  • z_dim - the size of the vector from which the image generation begins.

  • BATCH_SIZE - batch size for the loader dataset.

  • UPD_FOR_GEN - the number of discriminator updates per generator update.

  • epochs - number of epochs for training.

  • Dataset - the path to the dataset with images.

  • Weight_dir - the directory where the weights will be saved after each epoch.

In addition to the required parameters, additional parameters can also be passed to non-standard Trainer, for example, the following parameter is needed for "gp": Loss_config = {"lambda_gp": 10}.