-
Notifications
You must be signed in to change notification settings - Fork 1
/
main.py
90 lines (68 loc) · 2.34 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
from utils import *
import torch
import torch.optim as optim
import cv2
from torchvision.utils import save_image
from argparse import ArgumentParser
from weights import *
from tqdm import tqdm
from time import time
# Argument Parsing
parser = ArgumentParser(
description="Content and style arguments for the neural style transfer model"
)
parser.add_argument(
"--content",
help="Content image for the neural style transfer model",
default="content_sample.jpg",
)
parser.add_argument(
"--style",
help="Style image for the neural style transfer model",
default="style_sample.jpg",
)
args = parser.parse_args()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def main():
"""
Main function for training the Neural Style Transfer model
"""
start = time()
content = load_image(f"./data/{args.content}").to(device)
style = load_image(f"./data/{args.style}", shape=content.shape[-2:]).to(
device
)
model = get_model()
content_features = extract_features(content, model)
style_features = extract_features(style, model)
style_grams = {
layer: gram_matrix(style_features[layer]) for layer in style_features
}
target = content.clone().requires_grad_(True).to(device)
optimizer = optim.Adam([target], lr=0.01)
steps = 20000
# Main training loop
print("Starting training run")
for step in tqdm(range(1, steps + 1)):
target_features = extract_features(target, model)
content_loss = torch.mean(
(target_features["conv4_2"] - content_features["conv4_2"]) ** 2
)
style_loss = 0
for layer in style_weights:
target_feature = target_features[layer]
_, d, h, w = target_feature.shape
target_gram = gram_matrix(target_feature)
style_gram = style_grams[layer]
layer_style_loss = style_weights[layer] * torch.mean(
(target_gram - style_gram) ** 2
)
style_loss += layer_style_loss / (d * h * w)
total_loss = content_loss * content_loss + style_weight * style_loss
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
save_image(target, "./data/out_img.jpg")
print(f"Run completed in {(time() - start) / 60 :.2f} minutes.")
if __name__ == "__main__":
main()