style loss

st05.py

def style_loss(s_grams, t_features, weights):
    """
    Compute style loss, i.e. the weighted sum of MSE of all layers. 
    """
    # for each style feature, get target and style gramians, compare
    loss = 0
    for layer in weights:
        _, d, h, w = s_features[layer].shape
        t_gram = gramian(t_features[layer])    
        
        layer_loss = torch.mean((t_gram - s_grams[layer]) ** 2) / (d * h * w)
        loss += layer_loss * weights[layer]
    
    return loss