ProGamerGov / neural-style-pt

PyTorch implementation of neural style transfer algorithm
MIT License
833 stars 178 forks source link

HalfTensor (float16) support? tried to implement but failed :( #97

Open arpitest opened 1 year ago

arpitest commented 1 year ago

Is it possible to run this with half precision, to be able to use higher image resolution with limited VRAM? I've tried to do it (similar way to how stable diffusion does):

Running optimization with L-BFGS Iteration 10 / 1000 Content 1 loss: nan Style 1 loss: nan Style 2 loss: nan Style 3 loss: nan Style 4 loss: nan Style 5 loss: nan Total loss: nan

any idea how to fix that?

ziye3001 commented 1 year ago

well, I guess something exceeds the maximal value of float16, which is a fairly small number 65500. For example, the Gram Matrix is approximately the sum of the feature^2, which easily gets very large. The MSE loss of two Gram Matrices could have even larger magnitude. A possible solution could be some alternative losses, e.g. BN loss, in MMD paper. But you have to adjust the style weight to get similar results with the ones with Gram Matrices.

JCBrouwer commented 1 year ago

Another option would be to find out what the max feature values are on average and then add a fixed scaling before calculating the gram matrix.

This way the magnitude of the squaring operation is essentially limited to 1. You'd have to make sure that the loss weights are also adjusted accordingly though.