Normalization Issue

[ali-izhar]

The loss values diverge very quickly for non-compatible content and style images, and in some cases leading to NaN values. By non-compatible, I refer to images that are not fit for style transfer due to lighting etc. issues.

Is there a normalization mechanism in place that can regulate the loss while maintaining the quality of the "transfer"? I tried the following - which is helping - but I'm not sure if hard-coding it like this is best practice.

Normalizing the gram matrix values:

def gram_matrix(x):
    if x.dim() == 3:
        x = x.unsqueeze(0)  # Add a batch dimension if missing
    
    b, ch, h, w = x.size()
    features = x.view(b, ch, w * h)
    G = features.bmm(features.transpose(1, 2))
    return G.div(ch * h * w * b)    # additional normalization by batch size

Normalize TV loss:

def total_variation(y):
    batch_size, channels, height, width = y.size()
    tv_height = torch.abs(y[:, :, :-1, :] - y[:, :, 1:, :]).sum()
    tv_width = torch.abs(y[:, :, :, :-1] - y[:, :, :, 1:]).sum()
    return (tv_height + tv_width) / (batch_size * channels * height * width)