| import torch | |
| from data_setup import gram_matrix | |
| def content_loss(target, content): | |
| loss = torch.mean((target - content) ** 2) | |
| return loss | |
| def style_loss(target_features, style_grams): | |
| loss = 0 | |
| for layer in target_features: | |
| target_f = target_features[layer] | |
| target_gram = gram_matrix(target_f) | |
| style_gram = style_grams[layer] | |
| b,c,h,w = target_f.shape | |
| layer_loss = 0.2 * torch.mean((target_gram - style_gram) ** 2) | |
| loss += layer_loss/(c*h*w) | |
| return loss | |
| def total_loss(content_loss, style_loss, alpha, beta): | |
| loss = alpha * content_loss + beta * style_loss | |
| return loss | |