I’m trying to replicate the original StyleGAN’s batch size schedule: 128, 128, 128, 64, 32, 16 as the progressive growing is applied. I know I can recreate the DataLoader when I want to switch, but I’m working inside an extant framework that makes that a clunky change to make.
I never did figure out how to use collate_fn here so instead, I’m initializing my DataLoader with a batch_size of 16, and in my training loop I collect and concatenate these batches until I reach the actual batch size I want at any given time. This only works because all the batch sizes are divisible by 16. I tried to do this in collate_fn at first, I thought maybe it received a generator and I could return a different generator, but that wasn’t the case.
I’m still interested to know how collate_fn can be used to yield variable batch sizes, maybe it would be cleaner than my solution.