diff --git a/FastStyleTransferNode.py b/FastStyleTransferNode.py index 2f5a02c..89962ab 100644 --- a/FastStyleTransferNode.py +++ b/FastStyleTransferNode.py @@ -1,50 +1,22 @@ """ -This node is a simple conversion of this repository into ComfyUI ecosystem: +These nodes are a simple conversion of these repositories into ComfyUI ecosystem: https://github.com/rrmina/fast-neural-style-pytorch.git +https://github.com/gordicaleksa/pytorch-neural-style-transfer.git -Some of the code is written by ChatGPT4-o """ import torch import torch.nn as nn -from torchvision import models import time import os import folder_paths import subprocess as sp import sys - - - -class VGG16(nn.Module): - def __init__(self, vgg_path=os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/vgg/vgg16-00b39a1b.pth"), train=False): - super(VGG16, self).__init__() - # Load VGG Skeleton, Pretrained Weights - vgg16_features = models.vgg16(pretrained=False) - vgg16_features.load_state_dict(torch.load(vgg_path), strict=False) - self.features = vgg16_features.features - - # Turn-off Gradient History (on for testing train) - for param in self.features.parameters(): - param.requires_grad = train - - def forward(self, x): - layers = {'3': 'relu1_2', '8': 'relu2_2', '15': 'relu3_3', '22': 'relu4_3'} - features = {} - for name, layer in self.features._modules.items(): - x = layer(x) - if name in layers: - features[layers[name]] = x - if (name=='22'): - break - - return features - - -class ConvLayer(nn.Module): +# ML classes +class ConvolutionalLayer(nn.Module): def __init__(self, in_channels, out_channels, kernel_size, stride, norm="instance"): - super(ConvLayer, self).__init__() + super(ConvolutionalLayer, self).__init__() # Padding Layers self.padding_size = kernel_size // 2 self.reflection_pad = nn.ReflectionPad2d(self.padding_size) @@ -77,9 +49,9 @@ class ResidualLayer(nn.Module): """ def __init__(self, channels=128, kernel_size=3): super(ResidualLayer, self).__init__() - self.conv1 = ConvLayer(channels, channels, kernel_size, stride=1) + self.conv1 = ConvolutionalLayer(channels, channels, kernel_size, stride=1) self.relu = nn.ReLU() - self.conv2 = ConvLayer(channels, channels, kernel_size, stride=1) + self.conv2 = ConvolutionalLayer(channels, channels, kernel_size, stride=1) def forward(self, x): identity = x # preserve residual @@ -89,9 +61,9 @@ class ResidualLayer(nn.Module): return out -class DeconvLayer(nn.Module): +class DeconvolutionalLayer(nn.Module): def __init__(self, in_channels, out_channels, kernel_size, stride, output_padding, norm="instance"): - super(DeconvLayer, self).__init__() + super(DeconvolutionalLayer, self).__init__() # Transposed Convolution padding_size = kernel_size // 2 @@ -121,11 +93,11 @@ class TransformerNetwork(nn.Module): def __init__(self): super(TransformerNetwork, self).__init__() self.ConvBlock = nn.Sequential( - ConvLayer(3, 32, 9, 1), + ConvolutionalLayer(3, 32, 9, 1), nn.ReLU(), - ConvLayer(32, 64, 3, 2), + ConvolutionalLayer(32, 64, 3, 2), nn.ReLU(), - ConvLayer(64, 128, 3, 2), + ConvolutionalLayer(64, 128, 3, 2), nn.ReLU() ) self.ResidualBlock = nn.Sequential( @@ -136,11 +108,11 @@ class TransformerNetwork(nn.Module): ResidualLayer(128, 3) ) self.DeconvBlock = nn.Sequential( - DeconvLayer(128, 64, 3, 2, 1), + DeconvolutionalLayer(128, 64, 3, 2, 1), nn.ReLU(), - DeconvLayer(64, 32, 3, 2, 1), + DeconvolutionalLayer(64, 32, 3, 2, 1), nn.ReLU(), - ConvLayer(32, 3, 9, 1, norm="None") + ConvolutionalLayer(32, 3, 9, 1, norm="None") ) def forward(self, x): @@ -150,25 +122,27 @@ class TransformerNetwork(nn.Module): return out +# Node classes class TrainFastStyleTransfer: def __init__(self): pass @classmethod def INPUT_TYPES(s): - input_dir = folder_paths.get_input_directory() - files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))] return { "required": { - "style_img": (sorted(files), {"image_upload": True}), + "style_img": ("IMAGE",), "seed": ("INT", {"default": 30, "min": 0, "max": 999999, "step": 1,}), "content_weight": ("INT", {"default": 14, "min": 1, "max": 128, "step": 1,}), "style_weight": ("INT", {"default": 50, "min": 1, "max": 128, "step": 1,}), - "batch_size": ("INT", {"default": 4, "min": 1, "max": 128, "step": 1,}), - "train_img_size": ("INT", {"default": 256, "min": 256, "max": 2048, "step": 1,}), - "learning_rate": ("FLOAT", {"default": 0.001, "min": 0.0001, "max": 0.1, "step": 0.0001}), + "tv_weight": ("FLOAT", {"default": 0.001, "min": 0.0, "max": 1.0, "step": 0.0000001}), + "batch_size": ("INT", {"default": 4, "min": 1, "max": 32, "step": 1,}), + "train_img_size": ("INT", {"default": 256, "min": 128, "max": 2048, "step": 1,}), + "learning_rate": ("FLOAT", {"default": 0.001, "min": 0.0001, "max": 100.0, "step": 0.0001}), "num_epochs": ("INT", {"default": 1, "min": 1, "max": 20, "step": 1,}), - "save_model_every": ("INT", {"default": 500, "min": 100, "max": 10000, "step": 1,}), + "save_model_every": ("INT", {"default": 500, "min": 10, "max": 10000, "step": 10,}), + "from_pretrained": ("INT", {"default": 0, "min": 0, "max": 1, "step": 1,}), + "model": ([file for file in os.listdir(os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/models/")) if file.endswith('.pth')], ), }, } @@ -179,14 +153,19 @@ class TrainFastStyleTransfer: CATEGORY = "Style Transfer" + def encode_tensor(self, tensor): + tensor = tensor.permute(0, 3, 1, 2).contiguous() # Convert to [batch_size, channels, height, width] + return tensor[:, [2, 1, 0], :, :] * 255 - def train(self, style_img, seed, batch_size, train_img_size, learning_rate, num_epochs, content_weight, style_weight, save_model_every): + + def train(self, style_img, seed, batch_size, train_img_size, learning_rate, num_epochs, content_weight, style_weight, tv_weight, save_model_every, from_pretrained, model): + temp_save_style_img = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/temp/") + "temp_save_content_img.pt" save_model_path = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/models/") dataset_path = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/dataset/") vgg_path = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/vgg/vgg16-00b39a1b.pth") save_image_path = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/output/") - style_image_path = folder_paths.get_annotated_filepath(style_img) train_path = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/train.py") + command = [ @@ -195,24 +174,31 @@ class TrainFastStyleTransfer: '--dataset_path', dataset_path, '--vgg_path', vgg_path, '--num_epochs', str(num_epochs), - '--style_image_path', style_image_path, + '--temp_save_style_img', temp_save_style_img, '--batch_size', str(batch_size), '--content_weight', str(content_weight), '--style_weight', str(style_weight), + '--tv_weight', str(tv_weight), '--adam_lr', str(learning_rate), '--save_model_path', save_model_path, '--save_image_path', save_image_path, '--save_model_every', str(save_model_every), - '--seed', str(seed) + '--seed', str(seed), + '--pretrained_model' ] + if from_pretrained: + command.append(model) + else: + command.append('none') + + + torch.save(self.encode_tensor(style_img), temp_save_style_img) + sp.run(command) return () - - - class FastStyleTransfer: def __init__(self): pass @@ -262,15 +248,81 @@ class FastStyleTransfer: return (image,) +class NeuralStyleTransfer: + + + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "content_img": ("IMAGE",), + "style_img": ("IMAGE",), + "content_weight": ("FLOAT", {"default": 1e5, "min": 1e3, "max": 1e6, "step": 1e3}), + "style_weight": ("FLOAT", {"default": 3e4, "min": 1e1, "max": 1e5, "step": 1e1}), + "tv_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1e1, "step": 0.1}), + "num_steps": ("INT", {"default": 100, "min": 10, "max": 10000, "step": 10}), + "learning_rate": ("FLOAT", {"default": 1.0, "min": 1e-4, "max": 1e3, "step": 0.1}), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "neural_style_transfer" + CATEGORY = "Style Transfer" + + def encode_tensor(self, tensor): + tensor = tensor.permute(0, 3, 1, 2).contiguous() # Convert to [batch_size, channels, height, width] + return tensor * 255 + + def decode_tensor(self, tensor): + tensor = tensor.permute(0, 2, 3, 1).contiguous() # Convert to [batch_size, height, width, channels] + return tensor / 255 + + + def neural_style_transfer(self, content_img, style_img, content_weight, style_weight, tv_weight, num_steps, learning_rate): + + neural_style_transfer_path = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/neural_style_transfer.py") + + temp_save_content_img = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/temp/") + "temp_save_content_img.pt" + temp_save_style_img = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/temp/") + "temp_save_style_img.pt" + + temp_load_final_img = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/temp/") + "temp_load_final_img.pt" + + torch.save(self.encode_tensor(content_img), temp_save_content_img) + torch.save(self.encode_tensor(style_img), temp_save_style_img) + + + command = [ + sys.executable, neural_style_transfer_path, + '--content_weight', str(content_weight), + '--style_weight', str(style_weight), + '--tv_weight', str(tv_weight), + '--temp_save_style_img', temp_save_style_img, + '--temp_save_content_img', temp_save_content_img, + '--temp_load_final_img', temp_load_final_img, + '--num_steps', str(num_steps), + '--learning_rate', str(learning_rate) + ] + + sp.run(command) + + image = self.decode_tensor(torch.load(temp_load_final_img)) + os.remove(temp_save_style_img) + os.remove(temp_save_content_img) + os.remove(temp_load_final_img) + return (image,) + # A dictionary that contains all nodes you want to export with their names # NOTE: names should be globally unique NODE_CLASS_MAPPINGS = { "FastStyleTransfer": FastStyleTransfer, - "TrainFastStyleTransfer": TrainFastStyleTransfer + "TrainFastStyleTransfer": TrainFastStyleTransfer, + "NeuralStyleTransfer": NeuralStyleTransfer, } # A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { "FastStyleTransfer": "Fast Style Transfer", - "TrainFastStyleTransfer": "Train Fast Style Transfer" + "TrainFastStyleTransfer": "Train Fast Style Transfer", + "NeuralStyleTransfer": "Neural Style Transfer", } diff --git a/neural_style_transfer.py b/neural_style_transfer.py new file mode 100644 index 0000000..6047a6d --- /dev/null +++ b/neural_style_transfer.py @@ -0,0 +1,167 @@ +import torch +from torch.optim import AdamW, LBFGS +from torchvision import models +import torch.nn as nn +import argparse +from collections import namedtuple + +class Vgg16(torch.nn.Module): + """Only those layers are exposed which have already proven to work nicely.""" + def __init__(self, requires_grad=False, show_progress=False): + super().__init__() + vgg_pretrained_features = models.vgg16(pretrained=True, progress=show_progress).features + self.layer_names = ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3'] + self.content_feature_maps_index = 1 # relu2_2 + self.style_feature_maps_indices = list(range(len(self.layer_names))) # all layers used for style representation + + self.slice1 = torch.nn.Sequential() + self.slice2 = torch.nn.Sequential() + self.slice3 = torch.nn.Sequential() + self.slice4 = torch.nn.Sequential() + for x in range(4): + self.slice1.add_module(str(x), vgg_pretrained_features[x]) + for x in range(4, 9): + self.slice2.add_module(str(x), vgg_pretrained_features[x]) + for x in range(9, 16): + self.slice3.add_module(str(x), vgg_pretrained_features[x]) + for x in range(16, 23): + self.slice4.add_module(str(x), vgg_pretrained_features[x]) + if not requires_grad: + for param in self.parameters(): + param.requires_grad = False + + def forward(self, x): + x = self.slice1(x) + relu1_2 = x + x = self.slice2(x) + relu2_2 = x + x = self.slice3(x) + relu3_3 = x + x = self.slice4(x) + relu4_3 = x + vgg_outputs = namedtuple("VggOutputs", self.layer_names) + out = vgg_outputs(relu1_2, relu2_2, relu3_3, relu4_3) + return out + + +def prepare_model(device): + model = Vgg16() + + content_feature_maps_index = model.content_feature_maps_index + style_feature_maps_indices = model.style_feature_maps_indices + layer_names = model.layer_names + + content_fms_index_name = (content_feature_maps_index, layer_names[content_feature_maps_index]) + style_fms_indices_names = (style_feature_maps_indices, layer_names) + return model.to(device).eval(), content_fms_index_name, style_fms_indices_names + + +def gram_matrix(x, should_normalize=True): + (b, ch, h, w) = x.size() + features = x.view(b, ch, w * h) + features_t = features.transpose(1, 2) + gram = features.bmm(features_t) + if should_normalize: + gram /= ch * h * w + return gram + + +def total_variation(y): + return torch.sum(torch.abs(y[:, :, :, :-1] - y[:, :, :, 1:])) + \ + torch.sum(torch.abs(y[:, :, :-1, :] - y[:, :, 1:, :])) + + +def build_loss(neural_net, optimizing_img, target_representations, content_feature_maps_index, style_feature_maps_indices, config): + target_content_representation = target_representations[0] + target_style_representation = target_representations[1] + + current_set_of_feature_maps = neural_net(optimizing_img) + + current_content_representation = current_set_of_feature_maps[content_feature_maps_index].squeeze(axis=0) + content_loss = torch.nn.MSELoss(reduction='mean')(target_content_representation, current_content_representation) + + style_loss = 0.0 + current_style_representation = [gram_matrix(x) for cnt, x in enumerate(current_set_of_feature_maps) if cnt in style_feature_maps_indices] + for gram_gt, gram_hat in zip(target_style_representation, current_style_representation): + style_loss += torch.nn.MSELoss(reduction='sum')(gram_gt[0], gram_hat[0]) + style_loss /= len(target_style_representation) + + tv_loss = total_variation(optimizing_img) + + total_loss = config['content_weight'] * content_loss + config['style_weight'] * style_loss + config['tv_weight'] * tv_loss + + return total_loss, content_loss, style_loss, tv_loss + +def make_tuning_step(neural_net, optimizer, target_representations, content_feature_maps_index, style_feature_maps_indices, config): + def tuning_step(optimizing_img): + optimizer.zero_grad() + total_loss, content_loss, style_loss, tv_loss = build_loss(neural_net, optimizing_img, target_representations, content_feature_maps_index, style_feature_maps_indices, config) + total_loss.backward() + optimizer.step() + return total_loss, content_loss, style_loss, tv_loss + return tuning_step + +def neural_style_transfer_from_tensors(content_img_tensor, style_img_tensor, config, num_steps, learning_rate): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + content_img = content_img_tensor.to(device) + style_img = style_img_tensor.to(device) + + optimizing_img = content_img.clone().requires_grad_(True) + + neural_net, content_feature_maps_index_name, style_feature_maps_indices_names = prepare_model(device) + + content_img_set_of_feature_maps = neural_net(content_img) + style_img_set_of_feature_maps = neural_net(style_img) + + target_content_representation = content_img_set_of_feature_maps[content_feature_maps_index_name[0]].squeeze(axis=0) + target_style_representation = [gram_matrix(x) for cnt, x in enumerate(style_img_set_of_feature_maps) if cnt in style_feature_maps_indices_names[0]] + target_representations = [target_content_representation, target_style_representation] + + optimizer = LBFGS((optimizing_img,), max_iter=num_steps, lr=learning_rate, line_search_fn='strong_wolfe') + cnt = 0 + + def closure(): + nonlocal cnt + if torch.is_grad_enabled(): + optimizer.zero_grad() + total_loss, content_loss, style_loss, tv_loss = build_loss(neural_net, optimizing_img, target_representations, content_feature_maps_index_name[0], style_feature_maps_indices_names[0], config) + if total_loss.requires_grad: + total_loss.backward() + if cnt%100==0: + with torch.no_grad(): + print(f'L-BFGS | iteration: {cnt:03}, total loss={total_loss.item():12.4f}, content_loss={config["content_weight"] * content_loss.item():12.4f}, style loss={config["style_weight"] * style_loss.item():12.4f}, tv loss={config["tv_weight"] * tv_loss.item():12.4f}') + + cnt += 1 + return total_loss + + optimizer.step(closure) + + return optimizing_img.detach() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--temp_save_style_img", type=str) + parser.add_argument("--temp_save_content_img", type=str) + parser.add_argument("--temp_load_final_img", type=str) + + parser.add_argument("--content_weight", type=float, help="weight factor for content loss", default=1e5) + parser.add_argument("--style_weight", type=float, help="weight factor for style loss", default=3e4) + parser.add_argument("--tv_weight", type=float, help="weight factor for total variation loss", default=1e0) + parser.add_argument("--learning_rate", type=float, help="learning_rate", default=1e1) + parser.add_argument("--num_steps", type=int, help="number of training steps", default=100) + + args = parser.parse_args() + config = { + "content_weight": args.content_weight, + "style_weight": args.style_weight, + "tv_weight": args.tv_weight, + } + style_img = torch.load(args.temp_save_style_img) + content_img = torch.load(args.temp_save_content_img) + final_img = neural_style_transfer_from_tensors(content_img, style_img, config, args.num_steps, args.learning_rate) + print("nst done") + + torch.save(final_img, args.temp_load_final_img) + diff --git a/temp_files_go_here b/temp_files_go_here new file mode 100644 index 0000000..e69de29 diff --git a/train.py b/train.py index a51bb7a..e193cf7 100644 --- a/train.py +++ b/train.py @@ -7,8 +7,6 @@ import numpy as np import time import argparse import cv2 -import matplotlib.pyplot as plt - parser = argparse.ArgumentParser(description="Train a neural network with style transfer.") @@ -16,15 +14,17 @@ parser.add_argument('--train_image_size', type=int, default=256, help='Size of t parser.add_argument('--dataset_path', type=str, default="dataset", help='Path to the dataset') parser.add_argument('--vgg_path', type=str, default="vgg", help='Path to the vgg model') parser.add_argument('--num_epochs', type=int, default=1, help='Number of epochs for training') -parser.add_argument('--style_image_path', type=str, default="images/oil.jpg", help='Path to the style image') +parser.add_argument('--temp_save_style_img', type=str, default="temp/temp_save_style_img.pt", help='Path to the style image tensor') parser.add_argument('--batch_size', type=int, default=12, help='Batch size for training') parser.add_argument('--content_weight', type=float, default=8, help='Weight for content loss') parser.add_argument('--style_weight', type=float, default=50, help='Weight for style loss') +parser.add_argument('--tv_weight', type=float, default=0.001, help='Weight for total variation loss') parser.add_argument('--adam_lr', type=float, default=0.001, help='Learning rate for Adam optimizer') parser.add_argument('--save_model_path', type=str, default="models/oil/", help='Path to save the trained model') parser.add_argument('--save_image_path', type=str, default="images/out/", help='Path to save the output images') parser.add_argument('--save_model_every', type=int, default=200, help='Save model every n batches') parser.add_argument('--seed', type=int, default=1234, help='Random seed') +parser.add_argument('--pretrained_model', type=str, default='none', help='pretrained model') args = parser.parse_args() @@ -32,89 +32,34 @@ args = parser.parse_args() TRAIN_IMAGE_SIZE = args.train_image_size DATASET_PATH = args.dataset_path NUM_EPOCHS = args.num_epochs -STYLE_IMAGE_PATH = args.style_image_path +STYLE_IMAGE_PATH = args.temp_save_style_img VGG_PATH = args.vgg_path BATCH_SIZE = args.batch_size CONTENT_WEIGHT = args.content_weight STYLE_WEIGHT = args.style_weight +TV_WEIGHT = args.tv_weight ADAM_LR = args.adam_lr SAVE_MODEL_PATH = args.save_model_path SAVE_IMAGE_PATH = args.save_image_path SAVE_MODEL_EVERY = args.save_model_every SEED = args.seed -PLOT_LOSS = 1 +MODEL = args.pretrained_model # Utils - -def plot_loss_hist(c_loss, s_loss, total_loss, title="Loss History"): - x = [i for i in range(len(total_loss))] - plt.figure(figsize=[10, 6]) - plt.plot(x, c_loss, label="Content Loss") - plt.plot(x, s_loss, label="Style Loss") - plt.plot(x, total_loss, label="Total Loss") - - plt.legend() - plt.xlabel('Every 500 iterations') - plt.ylabel('Loss') - plt.title(title) - plt.show() - # Gram Matrix def gram(tensor): B, C, H, W = tensor.shape x = tensor.view(B, C, H*W) x_t = x.transpose(1, 2) - return torch.bmm(x, x_t) / (C*H*W) - -# Load image file -def load_image(path): - # Images loaded as BGR - img = cv2.imread(path) - return img - -# Show image -def show(img): - # Convert from BGR to RGB - img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) - - # imshow() only accepts float [0,1] or int [0,255] - img = np.array(img/255).clip(0,1) - - plt.figure(figsize=(10, 5)) - plt.imshow(img) - plt.show() - + gram_matrix = torch.bmm(x, x_t) / (C*H*W) + gram_matrix = torch.clamp(gram_matrix, min=1e-6, max=1e6) + return gram_matrix +# Save image def saveimg(img, image_path): img = img.clip(0, 255) cv2.imwrite(image_path, img) -# Preprocessing ~ Image to Tensor -def itot(img, max_size=None): - # Rescale the image - if (max_size==None): - itot_t = transforms.Compose([ - #transforms.ToPILImage(), - transforms.ToTensor(), - transforms.Lambda(lambda x: x.mul(255)) - ]) - else: - H, W, C = img.shape - image_size = tuple([int((float(max_size) / max([H,W]))*x) for x in [H, W]]) - itot_t = transforms.Compose([ - transforms.ToPILImage(), - transforms.Resize(image_size), - transforms.ToTensor(), - transforms.Lambda(lambda x: x.mul(255)) - ]) - - # Convert image to tensor - tensor = itot_t(img) - - # Add the batch_size dimension - tensor = tensor.unsqueeze(dim=0) - return tensor - # Preprocessing ~ Tensor to Image def ttoi(tensor): # Add the means @@ -130,6 +75,11 @@ def ttoi(tensor): img = img.transpose(1, 2, 0) return img +# Alternative loss functions experiments +def total_variation(y): + return torch.sum(torch.abs(y[:, :, :, :-1] - y[:, :, :, 1:])) + \ + torch.sum(torch.abs(y[:, :, :-1, :] - y[:, :, 1:, :])) + # VGG class VGG16(nn.Module): def __init__(self, vgg_path="models/vgg16-00b39a1b.pth"): @@ -164,11 +114,11 @@ class TransformerNetworkClass(nn.Module): def __init__(self): super(TransformerNetworkClass, self).__init__() self.ConvBlock = nn.Sequential( - ConvLayer(3, 32, 9, 1), + ConvolutionalLayer(3, 32, 9, 1), nn.ReLU(), - ConvLayer(32, 64, 3, 2), + ConvolutionalLayer(32, 64, 3, 2), nn.ReLU(), - ConvLayer(64, 128, 3, 2), + ConvolutionalLayer(64, 128, 3, 2), nn.ReLU() ) self.ResidualBlock = nn.Sequential( @@ -179,11 +129,11 @@ class TransformerNetworkClass(nn.Module): ResidualLayer(128, 3) ) self.DeconvBlock = nn.Sequential( - DeconvLayer(128, 64, 3, 2, 1), + DeconvolutionalLayer(128, 64, 3, 2, 1), nn.ReLU(), - DeconvLayer(64, 32, 3, 2, 1), + DeconvolutionalLayer(64, 32, 3, 2, 1), nn.ReLU(), - ConvLayer(32, 3, 9, 1, norm="None") + ConvolutionalLayer(32, 3, 9, 1, norm="None") ) def forward(self, x): @@ -192,9 +142,9 @@ class TransformerNetworkClass(nn.Module): out = self.DeconvBlock(x) return out -class ConvLayer(nn.Module): +class ConvolutionalLayer(nn.Module): def __init__(self, in_channels, out_channels, kernel_size, stride, norm="instance"): - super(ConvLayer, self).__init__() + super(ConvolutionalLayer, self).__init__() # Padding Layers padding_size = kernel_size // 2 self.reflection_pad = nn.ReflectionPad2d(padding_size) @@ -204,6 +154,7 @@ class ConvLayer(nn.Module): # Normalization Layers self.norm_type = norm + if (norm=="instance"): self.norm_layer = nn.InstanceNorm2d(out_channels, affine=True) elif (norm=="batch"): @@ -226,9 +177,9 @@ class ResidualLayer(nn.Module): """ def __init__(self, channels=128, kernel_size=3): super(ResidualLayer, self).__init__() - self.conv1 = ConvLayer(channels, channels, kernel_size, stride=1) + self.conv1 = ConvolutionalLayer(channels, channels, kernel_size, stride=1) self.relu = nn.ReLU() - self.conv2 = ConvLayer(channels, channels, kernel_size, stride=1) + self.conv2 = ConvolutionalLayer(channels, channels, kernel_size, stride=1) def forward(self, x): identity = x # preserve residual @@ -237,9 +188,9 @@ class ResidualLayer(nn.Module): out = out + identity # add residual return out -class DeconvLayer(nn.Module): +class DeconvolutionalLayer(nn.Module): def __init__(self, in_channels, out_channels, kernel_size, stride, output_padding, norm="instance"): - super(DeconvLayer, self).__init__() + super(DeconvolutionalLayer, self).__init__() # Transposed Convolution padding_size = kernel_size // 2 @@ -247,6 +198,7 @@ class DeconvLayer(nn.Module): # Normalization Layers self.norm_type = norm + if (norm=="instance"): self.norm_layer = nn.InstanceNorm2d(out_channels, affine=True) elif (norm=="batch"): @@ -261,8 +213,8 @@ class DeconvLayer(nn.Module): return out - def train(): + # Seeds torch.manual_seed(SEED) torch.cuda.manual_seed(SEED) @@ -286,15 +238,16 @@ def train(): # Load networks TransformerNetwork = TransformerNetworkClass().to(device) + if MODEL != 'none': + model_path = SAVE_MODEL_PATH + MODEL + print(f"Loading {model_path} model") + TransformerNetwork.load_state_dict(torch.load(model_path, map_location=device)) + TransformerNetwork.to(device) VGG = VGG16(vgg_path=VGG_PATH).to(device) # Get Style Features imagenet_neg_mean = torch.tensor([-103.939, -116.779, -123.68], dtype=torch.float32).reshape(1,3,1,1).to(device) - print(f"Train.py: style image path: {STYLE_IMAGE_PATH}") - style_image = load_image(STYLE_IMAGE_PATH) - print(f"Train.py: style image type: {type(style_image)}") - style_tensor = itot(style_image).to(device) - print(f"Train.py: style tensor type: {type(style_tensor)}") + style_tensor = torch.load(STYLE_IMAGE_PATH).to(device) style_tensor = style_tensor.add(imagenet_neg_mean) B, C, H, W = style_tensor.shape style_features = VGG(style_tensor.expand([BATCH_SIZE, C, H, W])) @@ -304,7 +257,7 @@ def train(): # Optimizer settings optimizer = optim.AdamW(TransformerNetwork.parameters(), lr=ADAM_LR, fused=True, amsgrad=True) - + scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=20, T_mult=2, eta_min=1e-6, last_epoch=-1) # Loss trackers content_loss_history = [] style_loss_history = [] @@ -312,22 +265,13 @@ def train(): batch_content_loss_sum = 0 batch_style_loss_sum = 0 batch_total_loss_sum = 0 - # Optimization/Training Loop - batch_count = 1 start_time = time.time() for epoch in range(NUM_EPOCHS): print("========Epoch {}/{}========".format(epoch+1, NUM_EPOCHS)) - for content_batch, _ in train_loader: - # Get current batch size in case of odd batch sizes - curr_batch_size = content_batch.shape[0] - - # Free-up unneeded cuda memory - torch.cuda.empty_cache() - - # Zero-out Gradients + for batch_count, (content_batch, _) in enumerate(train_loader): optimizer.zero_grad() - + torch.cuda.empty_cache() # Generate images and get features content_batch = content_batch[:,[2,1,0]].to(device) generated_batch = TransformerNetwork(content_batch) @@ -336,55 +280,63 @@ def train(): # Content Loss MSELoss = nn.MSELoss().to(device) - content_loss = CONTENT_WEIGHT * MSELoss(generated_features['relu2_2'], content_features['relu2_2']) + MAELoss = nn.L1Loss().to(device) + content_loss = CONTENT_WEIGHT * MSELoss(generated_features['relu2_2'], content_features['relu2_2']) + MAELoss(generated_features['relu2_2'], content_features['relu2_2']) / 2 batch_content_loss_sum += content_loss # Style Loss - style_loss = 0 + style_loss = 0.0 for key, value in generated_features.items(): - s_loss = MSELoss(gram(value), style_gram[key][:curr_batch_size]) - style_loss += s_loss - style_loss *= STYLE_WEIGHT + style_loss += MSELoss(gram(value), style_gram[key][:content_batch.shape[0]]) + style_loss += MAELoss(gram(value), style_gram[key][:content_batch.shape[0]]) + style_loss *= STYLE_WEIGHT/2 batch_style_loss_sum += style_loss.item() - # Total Loss - total_loss = content_loss + style_loss - batch_total_loss_sum += total_loss.item() + # TV loss + tv_loss = total_variation(generated_batch) * TV_WEIGHT - # Backprop and Weight Update + # Total Loss + total_loss = content_loss + style_loss + tv_loss + batch_total_loss_sum += total_loss.item() + total_loss.backward() optimizer.step() + if batch_count % 50 == 0: + scheduler.step() + + with torch.no_grad(): + print(f'AdamW | iteration: {batch_count+1:03}, total loss={total_loss.item():12.4f}, content_loss={content_loss.item():12.4f}, style loss={style_loss.item():12.4f}, tv loss={tv_loss.item():12.4f}') # Save Model and Print Losses - if (((batch_count-1)%SAVE_MODEL_EVERY == 0) or (batch_count==NUM_EPOCHS*len(train_loader))): + if (((batch_count)%SAVE_MODEL_EVERY == 0) or ((batch_count+1)==NUM_EPOCHS*len(train_loader))): # Print Losses - print("========Iteration {}/{}========".format(batch_count, NUM_EPOCHS*len(train_loader))) - print("\tContent Loss:\t{:.2f}".format(batch_content_loss_sum/batch_count)) - print("\tStyle Loss:\t{:.2f}".format(batch_style_loss_sum/batch_count)) - print("\tTotal Loss:\t{:.2f}".format(batch_total_loss_sum/batch_count)) + print("========Iteration {}/{}========".format(batch_count+1, NUM_EPOCHS*len(train_loader))) + print("\tContent Loss:\t{:.2f}".format(batch_content_loss_sum/(batch_count+1))) + print("\tStyle Loss:\t{:.2f}".format(batch_style_loss_sum/(batch_count+1))) + print("\tTotal Loss:\t{:.2f}".format(batch_total_loss_sum/(batch_count+1))) print("Time elapsed:\t{} seconds".format(time.time()-start_time)) # Save Model - checkpoint_path = SAVE_MODEL_PATH + "checkpoint_" + str(batch_count-1) + ".pth" + checkpoint_path = SAVE_MODEL_PATH + "checkpoint_" + str((batch_count+1)) + ".pth" torch.save(TransformerNetwork.state_dict(), checkpoint_path) print("Saved TransformerNetwork checkpoint file at {}".format(checkpoint_path)) # Save sample generated image sample_tensor = generated_batch[0].clone().detach().unsqueeze(dim=0) sample_image = ttoi(sample_tensor.clone().detach()) - sample_image_path = SAVE_IMAGE_PATH + "sample0_" + str(batch_count-1) + ".png" + sample_image_path = SAVE_IMAGE_PATH + "sample0_" + str((batch_count+1)) + ".png" saveimg(sample_image, sample_image_path) print("Saved sample tranformed image at {}".format(sample_image_path)) # Save loss histories - content_loss_history.append(batch_content_loss_sum/batch_count) - style_loss_history.append(batch_style_loss_sum/batch_count) - total_loss_history.append(batch_total_loss_sum/batch_count) + content_loss_history.append(batch_content_loss_sum/(batch_count+1)) + style_loss_history.append(batch_style_loss_sum/(batch_count+1)) + total_loss_history.append(batch_total_loss_sum/(batch_count+1)) + + + + - # Iterate Batch Counter - batch_count+=1 - if batch_count%10 == 0: - print(f"Batch count is {batch_count}") stop_time = time.time() # Print loss histories @@ -397,17 +349,7 @@ def train(): print("========Total Loss========") print(total_loss_history) - # Save TransformerNetwork weights - TransformerNetwork.eval() - TransformerNetwork.cpu() - final_path = SAVE_MODEL_PATH + "transformer_weight.pth" - print("Saving TransformerNetwork weights at {}".format(final_path)) - torch.save(TransformerNetwork.state_dict(), final_path) - print("Done saving final model") - - # Plot Loss Histories - if (PLOT_LOSS): - plot_loss_hist(content_loss_history, style_loss_history, total_loss_history) if __name__ == "__main__": train() +