CycleGAN icon indicating copy to clipboard operation
CycleGAN copied to clipboard

Freezing instance normalization during test phase?

Open Quasimondo opened this issue 8 years ago • 3 comments

I am looking for a way to get more consistency in the output when processing a bigger image in tiles.

The problem is that due to the memory constraints currently it's hard training and processing inputs bigger than 512x512 (actually for 64 filters the maximum I can train is 256x256 on a GTX 1080ti). So the workaround is to train and process the images in tiles, but trying to maintain the texture details and scale of the original image size. Of course this comes with its own bag of issues since the model does never see the entire image and might learn different features.

Anyway, it kind of works at least for artistic purposes, but the problem is that when processing all the tiles of an image in test mode the results of neighboring tiles often differ a lot in color (unfortunately also in style, which will be much harder to solve).

Here's an example: j_14779896064_80459dfff7_o

Philip Isola mentioned in a tweet that there might be a way to freeze the instance normalization, which possibly might help with the consistency. So my questions are: How would one do that? Is this something that has aready to be done during the training phase or it possible to just do it in the test phase?

I suspect that I have to look at the InstanceNormalization class' running_mean and runng_var, but any clue how to go about this would be very welcome.

Quasimondo avatar May 24 '17 09:05 Quasimondo

From http://pytorch.org/docs/nn.html:

At evaluation time (.eval()), the default behaviour of the InstanceNorm module stays the same i.e. running mean/variance is NOT used for normalization. One can force using stored mean and variance with .train(False) method.

However, the results might not look so good since instance norm isn't really meant to be used this way -- this will induce a big gap between how you trained instance norm and how it is being used at test time.

If it doesn't look good, you could try retraining in a way that removes this gap. One way would be "virtual instance norm", where you normalize one image based on the running statistics from a held out set of other images (same as virtual batchnorm from this paper).

Good luck and I'll be interested to know if you have success!

phillipi avatar May 26 '17 22:05 phillipi

Hi, torch implementation of Instance Norm is easily hackable and can use stored mean/var instead of independent normalization at test time. I think you just need to set self.bn to eval mode in https://github.com/junyanz/CycleGAN/blob/master/util/InstanceNormalization.lua#L102 or somewhere in updateOutput function. Stored statistics worked quite well for stylization.

Ideally, you want to get statistics from an entire image, you can create a dataset of tiles and compute global statistcs. I saw a script by @szagoruyko which does it, but could not find it now.

DmitryUlyanov avatar May 29 '17 06:05 DmitryUlyanov

@DmitryUlyanov Thanks for the notes.

@Quasimondo I guess the tiling operation will still cause the boundary artifacts even without instanceNorm.

Let's say your images are 512x512. I would recommend that you train your model on 256x256 cropped images. During the test time, you are still able to test it on large images (512x512), as testing requires less GPU memory, compared to training. In this case, you might want to set loadSize=512 and fineSize=256 for both training and test. Even you train a model on 256x256 images, you can still apply it fully convolutionally to a picture with arbitrary size.

See issue #19 for more discussion.

junyanz avatar Jun 13 '17 07:06 junyanz