nazianafis / neural-style-transfer Goto Github PK

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

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)