create app

.gitattributes

@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.pth.tar filter=lfs diff=lfs merge=lfs -text

adain.py

@@ -0,0 +1,37 @@
+import torch
+from utils import *
+
+
+class AdaIN(torch.nn.Module):
+    def __init__(self):
+        super(AdaIN, self).__init__()
+        # initialize instance normalization function
+        # this is the basis of our AdaIN layer, it follows an equation similar to a z-score
+        # (x - mu)/sigma
+        self.instance_norm = torch.nn.InstanceNorm2d(3)
+    
+    # forward method for our layer
+    # x would be the content input and y would be the style input
+    # both x and y are tensors
+    def forward(self, x, y):
+        # size is shaped (N, num_channels, Height, Width)
+        x_size = x.size()
+        
+        # we do not need these since they will be calculated by the instance normalization function
+        #x_mean, x_std = mean_and_std_of_image(x)
+        y_mean, y_std = mean_and_std_of_image(y)
+
+        x_norm = self.instance_norm(x)
+
+        
+        print(x_norm.size())
+        # expand size of tensors so that there are no shape errors when performing AdaIN operation
+        # if not self.training:
+        #     x_norm = x_norm.view(*x_norm.shape, 1)
+
+        x_size = x_norm.size()
+        print(x_size)
+        return y_std.expand(x_size) * x_norm + y_mean.expand(x_size)
+
+
+

app.py

@@ -0,0 +1,46 @@
+import gradio as gr
+import torch
+import torchvision.transforms as transforms
+import torch.nn as nn
+from decoder import decoder as Decoder
+from encoder import encoder as Encoder
+from net import StyleTransfer
+from PIL import Image
+
+encoder = Encoder
+decoder = Decoder
+
+encoder.load_state_dict(torch.load("./vgg_normalised.pth"))
+encoder = nn.Sequential(*list(encoder.children())[:31])
+decoder.load_state_dict(torch.load("./saved-models/decoder_iter_1000.pth.tar"))
+
+
+net = StyleTransfer(encoder, decoder)
+
+net.eval()
+
+def train_transform():
+        transform_list = [
+            transforms.Resize(size=(512, 512)),
+            # transforms.CenterCrop(256),
+            transforms.ToTensor()
+        ]
+        return transforms.Compose(transform_list)
+
+def cleanup(input, style):
+    transform = train_transform()
+    input_img = transform(Image.fromarray(input))
+    style_img = transform(Image.fromarray(style))
+    input_img = input_img.view(1, *input_img.shape)
+    style_img = style_img.view(1, *style_img.shape)
+    final_image_tensor = net(input_img, style_img)
+    final_image_tensor = final_image_tensor.squeeze()
+    to_pil = transforms.ToPILImage()
+    image = to_pil(final_image_tensor)
+    return image
+
+def greet(name):
+    return "Hello " + name + "!"
+
+demo = gr.Interface(fn=cleanup, inputs=[gr.Image(shape=(224, 224)),gr.Image(shape=(224,224))],outputs="image")
+demo.launch()

decoder.py

@@ -0,0 +1,33 @@
+import torch.nn as nn
+
+decoder = nn.Sequential(
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 256, (3, 3)),
+    nn.ReLU(),
+    nn.Upsample(scale_factor=2, mode='nearest'),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 128, (3, 3)),
+    nn.ReLU(),
+    nn.Upsample(scale_factor=2, mode='nearest'),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(128, 128, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(128, 64, (3, 3)),
+    nn.ReLU(),
+    nn.Upsample(scale_factor=2, mode='nearest'),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(64, 64, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(64, 3, (3, 3)),
+)

encoder.py

@@ -0,0 +1,58 @@
+import torch.nn as nn
+
+
+encoder = nn.Sequential(
+    nn.Conv2d(3, 3, (1, 1)),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(3, 64, (3, 3)),
+    nn.ReLU(),  # relu1-1
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(64, 64, (3, 3)),
+    nn.ReLU(),  # relu1-2
+    nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(64, 128, (3, 3)),
+    nn.ReLU(),  # relu2-1
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(128, 128, (3, 3)),
+    nn.ReLU(),  # relu2-2
+    nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(128, 256, (3, 3)),
+    nn.ReLU(),  # relu3-1
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),  # relu3-2
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),  # relu3-3
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),  # relu3-4
+    nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 512, (3, 3)),
+    nn.ReLU(),  # relu4-1, this is the last layer used
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu4-2
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu4-3
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu4-4
+    nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu5-1
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu5-2
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu5-3
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU()  # relu5-4
+)

logs/events.out.tfevents.1673075465.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dce644727440026ca34adf6e356baa976b1abdad898c4227947623aaa6c27242
+size 40

logs/events.out.tfevents.1673075531.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9ff3282e663a03174d06e037b601a40d8796d6ac3c2f02dc803ae854bd47f224
+size 40

logs/events.out.tfevents.1673075820.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82618836bcd68966be78e10fecbf7316cca493eb6131a82dfcbbcfcd1d5ec66a
+size 40

logs/events.out.tfevents.1673075821.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ccbaca8977b8c54996f333c9cd844a3b4cd6f81f619f4e2cd355d96846110c2f
+size 40

logs/events.out.tfevents.1673075850.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ac12b47b00399439a60549dafeceba4525b0fbabf7c1046c63d2559fdd719562
+size 40

logs/events.out.tfevents.1673075852.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f1e5c0d77d3b9ed00d7bba6febcb5568f4ea1d01c77a13b48e5fa7c3bfe088b4
+size 40

logs/events.out.tfevents.1673075889.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f5ade4d30a4b1a32752bfe04ccf7499bf56fbb82d279ccc024254f465be3b253
+size 40

logs/events.out.tfevents.1673075890.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:826d5080380cadac7bd1525880ef05aab8be1ed4c9ccf5c1d7349f1c6280ca03
+size 40

logs/events.out.tfevents.1673075982.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f733d66ed2181abd3f250e23025231ee530cab128f976c282b502706834b11af
+size 40

logs/events.out.tfevents.1673076026.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5bc909a1adc9130a19913a9e53e634167b20cb15c50c8cc35d3c6d9a92b2f146
+size 40

logs/events.out.tfevents.1673076079.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b03cba18d7045162a22c08091d1492af2ef99cd211a562f5bc82a1752be8c3c9
+size 40

logs/events.out.tfevents.1673076142.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:156bf3ee127ca32ca84fa3405acdc4ffc1431a9dd96d5ccd5c5267cb05aaed62
+size 40

logs/events.out.tfevents.1673076233.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e687762ebd774a4b8b84609559927223587f4f59128e41360d6e8b3c77d21320
+size 40

logs/events.out.tfevents.1673076507.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c07751136a8cea75e817528ee97e0d018c9779150bbd69a431ca7aa35aeb49cd
+size 40

logs/events.out.tfevents.1673076723.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bac6375530aa11a872827fd9006f07f37f9122c475b6398f4afbced98fa575d4
+size 40

logs/events.out.tfevents.1673076832.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f82096b51b53c03476c219d7cef64374f21277e4f7d11a98877786458e215f1f
+size 40

logs/events.out.tfevents.1673076887.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a33b760c94457ccf05b3b8c81c4c7a0dd2af318c72dafbb94388fb805e99cc0
+size 40

logs/events.out.tfevents.1673076993.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1783ef3a87692e3d75c763073e30c9c5d9f0f577017cff2c3cee63a5e565dbf2
+size 40

logs/events.out.tfevents.1673077155.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d5baedc0a372f0ac969de671e31c15a60e7185972a15fc730cda35add60c51a9
+size 40

logs/events.out.tfevents.1673077187.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5293293201e7252e90e863e550f2facdf95c7a84970424003f44e2e9285e0d0e
+size 40

logs/events.out.tfevents.1673077234.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c998d2738fb5d91c2328fc55968fc59cb021c485c94b0e0eec33d043cd5b92f4
+size 40

logs/events.out.tfevents.1673079573.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:439a6b8f06eacebeade084bc0b99cdf0fc2dd664f254f63035b85e6bea864954
+size 40

logs/events.out.tfevents.1673079783.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:19fa8348ce1a15cbcd2fbce80cdc4ca5935e2a776975dc914120b8e1989a77e5
+size 40

logs/events.out.tfevents.1673079809.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dd2327969e02ac17d9fd678ecb993897c2d48d1257cdb60043a7bfeed7188db3
+size 40

logs/events.out.tfevents.1673079875.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:07c018ca3aed8cc142059b3c38a35ce270db39d4e830bec58e2d432a6539ce97
+size 40

logs/events.out.tfevents.1673079932.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5e140d7be63757c53ae28febf5a730b453986307698009cb0c1163cd3adabed
+size 40

logs/events.out.tfevents.1673080014.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:646d967ffd7850501fb6bdde7df8aa44850d27f8d1e9011816b25531116d06f5
+size 40

logs/events.out.tfevents.1673080084.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08baf672c0d6ab6eb18d6c0e4dd4b45d061171e346b51159e034497ccea12123
+size 44786

logs/events.out.tfevents.1673080471.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5a7d92ffc0a3f73e0a861c3b6cd709e9638e904442c6115459ce7faceced9d33
+size 4940

logs/events.out.tfevents.1673080709.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0b90008548a1bbd7852e46eb4b71dcf412f634bf93be55d7b84bce046cbe8fc6
+size 109386

logs/events.out.tfevents.1673733387.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8fd458f41aa7271947c88f846db26cb0a58698505d219791ae5814accad1bab
+size 109386

logs/events.out.tfevents.1673735400.Vikrams-MBP.lan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb80a6f53db3f15fc30d4cf0267b4ce6b2b87288ae374949db56ce4d57fcb07b
+size 109386

net.py

@@ -0,0 +1,76 @@
+import torch
+import torch.nn as nn
+from adain import AdaIN
+from utils import *
+
+class StyleTransfer(nn.Module):
+    def __init__(self, encoder, decoder):
+        super(StyleTransfer, self).__init__()
+        layers = list(encoder.children())
+        self.enc_1 = nn.Sequential(*layers[:4])  # input -> relu1_1
+        self.enc_2 = nn.Sequential(*layers[4:11])  # relu1_1 -> relu2_1
+        self.enc_3 = nn.Sequential(*layers[11:18])  # relu2_1 -> relu3_1
+        self.enc_4 = nn.Sequential(*layers[18:31])  # relu3_1 -> relu4_1]
+        self.relus = [self.enc_1, self.enc_2, self.enc_3, self.enc_4]
+        self.decoder = decoder
+        self.mse = nn.MSELoss()
+        self.adain = AdaIN()
+
+        for name in ['enc_1', 'enc_2', 'enc_3', 'enc_4']:
+            for param in getattr(self, name).parameters():
+                param.requires_grad = False
+
+    def encode_with_save(self, input):
+        results = [input]
+        for i in range(4):
+            func = getattr(self, 'enc_{:d}'.format(i + 1))
+            results.append(func(results[-1]))
+        return results[1:]
+        
+    def encode(self, input):
+        res = input
+        for layer in self.relus:
+            res = layer(res)
+        return res
+    
+    def forward(self, content, style):
+        if not self.training:
+            self.adain.eval()
+        encoded_style = self.encode_with_save(style)
+        encoded_content = self.encode(content)
+        t = self.adain(encoded_content, encoded_style[-1])
+
+
+        g_t = self.decoder(t)
+
+        if not self.training:
+            return g_t
+        g_t_encoding = self.encode_with_save(g_t)
+
+        s_loss = self.style_loss(g_t_encoding, encoded_style)
+        c_loss = self.content_loss(g_t_encoding[-1], t)
+
+        return g_t, s_loss, c_loss
+
+
+    def style_loss(self, encoded_image, encoded_style):
+        MSE = torch.nn.MSELoss()
+        initial_mean_image, initial_std_image = mean_and_std_of_image(encoded_image[0])
+        initial_mean_style, initial_std_style = mean_and_std_of_image(encoded_style[0])
+        loss = MSE(initial_mean_image, initial_mean_style) + MSE(initial_std_image, initial_std_style)
+        for i in range(1, 4, 1):
+            mean_image, std_image = mean_and_std_of_image(encoded_image[i])
+            mean_style, std_style = mean_and_std_of_image(encoded_style[i])
+            loss += MSE(mean_image, mean_style) + MSE(std_image, std_style)
+        return loss
+
+
+    def content_loss(self, encoded_image, style_content_combined):
+        MSE = torch.nn.MSELoss()
+        return MSE(encoded_image, style_content_combined)
+
+
+
+
+
+        

saved-models/decoder_iter_1000.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5831e836f6bb7bab64fef12a1a070d69efb6ec312dc3c4a96653e22af55b5809
+size 14026951

saved-models/decoder_iter_500.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08bd329dfe5e930bccebf1045365818574eb203008c9ea3a30fe53c657ac6b32
+size 14026931

train.py

@@ -0,0 +1,144 @@
+from net import StyleTransfer
+import torch
+import torch.nn as nn
+from pathlib import Path
+import torchvision
+import torch.utils.data as data
+import torchvision.transforms as transforms
+import matplotlib.pyplot as plt
+import torch.multiprocessing
+from utils import *
+import argparse
+from tqdm import tqdm
+from tensorboardX import SummaryWriter
+from decoder import decoder as Decoder
+from encoder import encoder as Encoder
+from PIL import Image, ImageFile
+
+class FlatFolderDataset(data.Dataset):
+        def __init__(self, root, transform):
+            super(FlatFolderDataset, self).__init__()
+            self.root = root
+            self.paths = list(Path(self.root).glob('*'))
+            self.transform = transform
+
+        def __getitem__(self, index):
+            path = self.paths[index]
+            img = Image.open(str(path)).convert('RGB')
+            img = self.transform(img)
+            return img
+
+        def __len__(self):
+            return len(self.paths)
+
+        def name(self):
+            return 'FlatFolderDataset'
+
+def main():
+    torch.multiprocessing.set_sharing_strategy('file_system')
+
+    # Set the path to the dataset directory
+    content_dataset_dir = '../../content-dataset/images/images'
+    style_dataset_dir = '../../style-dataset/images'
+
+
+    def train_transform():
+        transform_list = [
+            transforms.Resize(size=(512, 512)),
+            transforms.RandomCrop(256),
+            transforms.ToTensor()
+        ]
+        return transforms.Compose(transform_list)
+
+
+    
+
+    parser = argparse.ArgumentParser()
+    # Basic options
+    parser.add_argument('--content_dir', default=content_dataset_dir, type=str,
+                        help='Directory path to a batch of content images')
+    parser.add_argument('--style_dir', default=style_dataset_dir, type=str,
+                        help='Directory path to a batch of style images')
+    parser.add_argument('--encoder', type=str, default='./vgg_normalised.pth')
+
+    # training options
+    parser.add_argument('--save_dir', default='../saved-models',
+                        help='Directory to save the model')
+    parser.add_argument('--log_dir', default='./logs',
+                        help='Directory to save the log')
+    parser.add_argument('--lr', type=float, default=1e-4)
+    parser.add_argument('--lr_decay', type=float, default=5e-5)
+    parser.add_argument('--max_iter', type=int, default=8000)
+    parser.add_argument('--batch_size', type=int, default=8)
+    parser.add_argument('--style_weight', type=float, default=10.0)
+    parser.add_argument('--content_weight', type=float, default=1.0)
+    parser.add_argument('--n_threads', type=int, default=8)
+    parser.add_argument('--save_model_interval', type=int, default=500)
+    parser.add_argument('--save-image-interval', type=int, default=50)
+    args = parser.parse_args()
+
+
+
+
+    device = torch.device('mps')
+    save_dir = Path(args.save_dir)
+    save_dir.mkdir(exist_ok=True, parents=True)
+    log_dir = Path(args.log_dir)
+    log_dir.mkdir(exist_ok=True, parents=True)
+    writer = SummaryWriter(log_dir=str(log_dir))
+
+
+    decoder = Decoder
+    encoder = Encoder
+
+    encoder.load_state_dict(torch.load(args.encoder))
+    encoder = nn.Sequential(*list(encoder.children())[:31])
+    network = StyleTransfer(encoder, decoder)
+    network.train()
+    network.to(device)
+
+    content_dataset = FlatFolderDataset(args.content_dir, transform=train_transform())
+    style_dataset = FlatFolderDataset(args.style_dir, transform=train_transform())
+
+    print(len(content_dataset), len(style_dataset))
+
+    content_iter = iter(data.DataLoader(
+        content_dataset, batch_size=args.batch_size,
+        num_workers=args.n_threads))
+    style_iter = iter(data.DataLoader(
+        style_dataset, batch_size=args.batch_size,
+        num_workers=args.n_threads))
+    optimizer = torch.optim.Adam(network.decoder.parameters(), lr=args.lr)
+
+
+    for batch in tqdm(range(args.max_iter)):
+        adjust_learning_rate(optimizer, batch, args.lr_decay, args.lr)
+        content_images = next(content_iter).to(device)
+        style_images = next(style_iter).to(device)
+        final_image, s_loss, c_loss = network(content_images, style_images)
+        c_loss = args.content_weight * c_loss
+        s_loss = args.style_weight * s_loss
+        total_loss = c_loss + s_loss
+
+        optimizer.zero_grad()
+        total_loss.backward()
+        optimizer.step()
+
+        writer.add_scalar('loss_content', c_loss.item(), batch + 1)
+        writer.add_scalar('loss_style', s_loss.item(), batch + 1)
+
+        if (batch + 1) % args.save_model_interval == 0 or (batch + 1) == args.max_iter:
+            state_dict = network.decoder.state_dict()
+            for key in state_dict.keys():
+                state_dict[key] = state_dict[key].to(torch.device('cpu'))
+            torch.save(state_dict, save_dir /
+                    'decoder_iter_{:d}.pth.tar'.format(batch + 1))
+
+        if (batch + 1) % args.save_image_interval == 0:
+            print_img = torch.cat((content_images[:1], style_images[:1], final_image[:1]), 3).detach().cpu()
+            concat_img(print_img, batch)
+    writer.close()
+
+
+if __name__ == "__main__":
+    main()

utils.py

@@ -0,0 +1,42 @@
+import torch
+import matplotlib.pyplot as plt
+import numpy as np
+
+def adjust_learning_rate(optimiser, iters, learning_rate_decay, LR):
+    for param_group in optimiser.param_groups:
+        param_group['lr'] = LR / (1.0 + learning_rate_decay * iters)
+
+def concat_img(imgs, batch):
+    plt.figure()
+    #imgs = (imgs + 1) / 2
+    imgs = imgs.movedim((0, 1, 2, 3), (0, 3, 1, 2)).detach().cpu().numpy() 
+    axs = plt.imshow(np.concatenate(imgs.tolist(), axis=1))
+    plt.axis('off')
+    plt.savefig("../../produced-images/batch{}img.png".format(batch))
+    plt.close()
+
+def concat_img(imgs, batch):
+    plt.figure()
+    #imgs = (imgs + 1) / 2
+    imgs = imgs.movedim((0, 1, 2, 3), (0, 3, 1, 2)).detach().cpu().numpy() 
+    axs = plt.imshow(np.concatenate(imgs.tolist(), axis=1))
+    plt.axis('off')
+    plt.savefig("../../produced-images/batch{}img.png".format(batch))
+
+# takes in image tensor x as input
+def mean_and_std_of_image(x):
+    x_size = x.size()
+    # turn x into the shape of (batch_size, num_channels, height*width)
+    x = x.view(x.shape[0], x.shape[1], -1)
+    #calculate the mean of the second dimension, H*W
+    mean = x.mean(dim=2)
+    std = x.var(dim=2).sqrt()
+    #reshape mean and std to size (batch_size, num_channels, 1, 1)
+    #because mean and std are sort of a scalar quantity the last two dimensions are both 1
+    # mean = mean.view(mean.shape[0], mean.shape[1], 1, 1)
+    # std = std.view(std.shape[0], std.shape[1], 1, 1)
+
+    mean = mean.view(mean.shape[0], mean.shape[1], 1, 1)
+    std = std.view(std.shape[0], std.shape[1], 1, 1)
+
+    return (mean, std)

vgg_normalised.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:804ca2835ecf7539f0cd2a7ac3c18ce81e6f8468969ae7117ac0c148d286bb4a
+size 80102481