Can't use stylegan2 256x256 models

[NitayGitHub]

Describe the bug
I wanted to use transfer learning on 256x256 pkl and my data contains images of 256x256 yet I got this error.

Training options:
{
  "G_kwargs": {
    "class_name": "training.networks_stylegan2.Generator",
    "z_dim": 256,
    "w_dim": 256,
    "mapping_kwargs": {
      "num_layers": 8,
      "freeze_layers": 0,
      "freeze_embed": false
    },
    "channel_base": 32768,
    "channel_max": 256,
    "fused_modconv_default": "inference_only"
  },
  "D_kwargs": {
    "class_name": "training.networks_stylegan2.Discriminator",
    "block_kwargs": {
      "freeze_layers": 0
    },
    "mapping_kwargs": {},
    "epilogue_kwargs": {
      "mbstd_group_size": 4
    },
    "channel_base": 32768,
    "channel_max": 256
  },
  "G_opt_kwargs": {
    "class_name": "torch.optim.Adam",
    "betas": [
      0,
      0.99
    ],
    "eps": 1e-08,
    "lr": 0.002
  },
  "D_opt_kwargs": {
    "class_name": "torch.optim.Adam",
    "betas": [
      0,
      0.99
    ],
    "eps": 1e-08,
    "lr": 0.002
  },
  "loss_kwargs": {
    "class_name": "training.loss.StyleGAN2Loss",
    "r1_gamma": 16.0,
    "style_mixing_prob": 0.9,
    "pl_weight": 2,
    "pl_no_weight_grad": true,
    "blur_init_sigma": 0
  },
  "data_loader_kwargs": {
    "pin_memory": true,
    "prefetch_factor": 2,
    "num_workers": 3
  },
  "training_set_kwargs": {
    "class_name": "training.dataset.ImageFolderDataset",
    "path": "./datasets/FH.zip",
    "use_labels": false,
    "max_size": 4592,
    "xflip": false,
    "yflip": false,
    "resolution": 256,
    "random_seed": 0
  },
  "num_gpus": 2,
  "batch_size": 16,
  "batch_gpu": 8,
  "metrics": [],
  "total_kimg": 25000,
  "resume_kimg": 360,
  "kimg_per_tick": 4,
  "network_snapshot_ticks": 10,
  "image_snapshot_ticks": 10,
  "snap_res": "4k",
  "random_seed": 0,
  "ema_kimg": 5.0,
  "G_reg_interval": 4,
  "augment_kwargs": {
    "class_name": "training.augment.AugmentPipe",
    "xflip": 1,
    "rotate90": 1,
    "xint": 1,
    "scale": 1,
    "rotate": 1,
    "aniso": 1,
    "xfrac": 1,
    "brightness": 1,
    "contrast": 1,
    "lumaflip": 1,
    "hue": 1,
    "saturation": 1
  },
  "ada_target": 0.6,
  "resume_pkl": "https://nvlabs-fi-cdn.nvidia.com/stylegan2/networks/stylegan2-cat-config-f.pkl",
  "ada_kimg": 100,
  "ema_rampup": null,
  "run_dir": "./results/00004-stylegan2-FH-gpus2-batch16-gamma16-resume_lsuncat256"
}

Output directory:    ./results/00004-stylegan2-FH-gpus2-batch16-gamma16-resume_lsuncat256
Number of GPUs:      2
Batch size:          16 images
Training duration:   25000 kimg
Dataset path:        ./datasets/FH.zip
Dataset size:        4592 images
Dataset resolution:  256
Dataset labels:      False
Dataset x-flips:     False
Dataset y-flips:     False

Launching processes...
Loading training set...

Num images:  4592
Image shape: [3, 256, 256]
Label shape: [0]
Downloading https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-ffhq-256x256.pkl ... done
Traceback (most recent call last):
  File "train.py", line 369, in <module>
    main()  # pylint: disable=no-value-for-parameter
  File "/opt/conda/lib/python3.7/site-packages/click/core.py", line 1128, in __call__
    return self.main(*args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/click/core.py", line 1053, in main
    rv = self.invoke(ctx)
  File "/opt/conda/lib/python3.7/site-packages/click/core.py", line 1395, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/opt/conda/lib/python3.7/site-packages/click/core.py", line 754, in invoke
    return __callback(*args, **kwargs)
  File "train.py", line 362, in main
    launch_training(c=c, desc=desc, outdir=opts.outdir, dry_run=opts.dry_run)
  File "train.py", line 94, in launch_training
    torch.multiprocessing.spawn(fn=subprocess_fn, args=(c, temp_dir), nprocs=c.num_gpus)
  File "/opt/conda/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 240, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/opt/conda/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 198, in start_processes
    while not context.join():
  File "/opt/conda/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 160, in join
    raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException: 

-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/opt/conda/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/kaggle/working/stylegan3-fun/train.py", line 50, in subprocess_fn
    training_loop.training_loop(rank=rank, **c)
  File "/kaggle/working/stylegan3-fun/training/training_loop.py", line 163, in training_loop
    misc.copy_params_and_buffers(resume_data[name], module, require_all=False)
  File "/kaggle/working/stylegan3-fun/torch_utils/misc.py", line 162, in copy_params_and_buffers
    tensor.copy_(src_tensors[name].detach()).requires_grad_(tensor.requires_grad)
RuntimeError: The size of tensor a (256) must match the size of tensor b (512) at non-singleton dimension 0

[nuclearsugar]

I believe that the model that is shown in the log above is 512x512. Instead you need to use a 256x256 model, such as:
https://api.ngc.nvidia.com/v2/models/org/nvidia/team/research/stylegan2/1/files?redirect=true&path=stylegan2-ffhq-256x256.pkl

Also when training a 256x256 model then be sure to include the following attribute in your training parameters.
--cbase=16384

[NitayGitHub]

I run
!python train.py --outdir=./results --cbase=16384 --snap=10 --img-snap=10 --cfg=stylegan2 --data=./datasets/FH.zip --augpipe=bgc --gpus=2 --metrics=None --gamma=12 --batch=16 --resume='https://api.ngc.nvidia.com/v2/models/org/nvidia/team/research/stylegan2/1/files?redirect=true&path=stylegan2-ffhq-256x256.pkl'

and got the same issue

[nuclearsugar]

I realize that I gave you that URL for the 256x256 model, but it's not a valid download link.

Try the code below (as documented here).

!python train.py --outdir=./results --cbase=16384 --snap=10 --img-snap=10 --cfg=stylegan2 --data=./datasets/FH.zip --augpipe=bgc --gpus=2 --metrics=None --gamma=12 --batch=16 --resume=ffhq256

[PDillis]

The error comes from the dimensionality in the latent space, as you have at the top of your configuration: "G_kwargs": {..., "z_dim": 256, "w_dim": 256, ...}. This is bizarre, as we set up the correct dimensionality here (and is the one that the pre-trained models are using). Perhaps there's somewhere else these values are being changed, but I'll have to look into it as the train.py file only changes this value for --cfg=stylegan2-ext.

[NitayGitHub]

I changed "z_dim" and "w_dim" to 256, thinking it might help but it didn't. However, I believe the problem was with the dataset I used where for some reason some images were not in 256x256 size. I added
if img.size != (256, 256):
img = img.resize((256, 256))

and it fixed it.
Although torchvision transforms.RandomCrop(size=256) should have made sure all images are in 256x256 resolution, it didn't.

[PDillis]

Yeah you need to exactly match the model you are finetuning from, otherwise there's no way to use the weights. For the reshaping of your data, do you mean you used dataset_tool.py and it still resulted in images of different size, or do you have another pipeline there?

[NitayGitHub]

Actually, it seems the reason it was fixed was thanks to adding --cbase=16384

[nuclearsugar]

Indeed, in my experience the --cbase=16384 is required when fine-tuning a 256x256 model. Otherwise it will throw an error prompt.