Add TrainFastStyleTransfer node
This commit is contained in:
zeroxoxo
+100
-6
@@ -5,19 +5,49 @@ https://github.com/rrmina/fast-neural-style-pytorch.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):
|
||||
def __init__(self, in_channels, out_channels, kernel_size, stride, norm="instance"):
|
||||
super(ConvLayer, self).__init__()
|
||||
# Padding Layers
|
||||
padding_size = kernel_size // 2
|
||||
self.reflection_pad = nn.ReflectionPad2d(padding_size)
|
||||
self.padding_size = kernel_size // 2
|
||||
self.reflection_pad = nn.ReflectionPad2d(self.padding_size)
|
||||
|
||||
# Convolution Layer
|
||||
self.conv_layer = nn.Conv2d(in_channels, out_channels, kernel_size, stride)
|
||||
@@ -32,7 +62,7 @@ class ConvLayer(nn.Module):
|
||||
def forward(self, x):
|
||||
x = self.reflection_pad(x)
|
||||
x = self.conv_layer(x)
|
||||
if (self.norm_type=="None"):
|
||||
if self.norm_type == "None":
|
||||
out = x
|
||||
else:
|
||||
out = self.norm_layer(x)
|
||||
@@ -120,6 +150,68 @@ class TransformerNetwork(nn.Module):
|
||||
return out
|
||||
|
||||
|
||||
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}),
|
||||
"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}),
|
||||
"num_epochs": ("INT", {"default": 1, "min": 1, "max": 20, "step": 1,}),
|
||||
"save_model_every": ("INT", {"default": 500, "min": 100, "max": 10000, "step": 1,}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
OUTPUT_NODE = True
|
||||
|
||||
FUNCTION = "train"
|
||||
|
||||
CATEGORY = "Style Transfer"
|
||||
|
||||
|
||||
def train(self, style_img, seed, batch_size, train_img_size, learning_rate, num_epochs, content_weight, style_weight, save_model_every):
|
||||
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 = [
|
||||
sys.executable, train_path,
|
||||
'--train_image_size', str(train_img_size),
|
||||
'--dataset_path', dataset_path,
|
||||
'--vgg_path', vgg_path,
|
||||
'--num_epochs', str(num_epochs),
|
||||
'--style_image_path', style_image_path,
|
||||
'--batch_size', str(batch_size),
|
||||
'--content_weight', str(content_weight),
|
||||
'--style_weight', str(style_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)
|
||||
]
|
||||
|
||||
sp.run(command)
|
||||
return ()
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
class FastStyleTransfer:
|
||||
def __init__(self):
|
||||
@@ -173,10 +265,12 @@ class FastStyleTransfer:
|
||||
# A dictionary that contains all nodes you want to export with their names
|
||||
# NOTE: names should be globally unique
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"FastStyleTransfer": FastStyleTransfer
|
||||
"FastStyleTransfer": FastStyleTransfer,
|
||||
"TrainFastStyleTransfer": TrainFastStyleTransfer
|
||||
}
|
||||
|
||||
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FastStyleTransfer": " Fast Style Transfer"
|
||||
"FastStyleTransfer": "Fast Style Transfer",
|
||||
"TrainFastStyleTransfer": "Train Fast Style Transfer"
|
||||
}
|
||||
|
||||
@@ -0,0 +1,413 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
from torchvision import datasets, transforms, models
|
||||
import random
|
||||
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.")
|
||||
|
||||
parser.add_argument('--train_image_size', type=int, default=256, help='Size of the training images')
|
||||
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('--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('--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')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# GLOBAL SETTINGS
|
||||
TRAIN_IMAGE_SIZE = args.train_image_size
|
||||
DATASET_PATH = args.dataset_path
|
||||
NUM_EPOCHS = args.num_epochs
|
||||
STYLE_IMAGE_PATH = args.style_image_path
|
||||
VGG_PATH = args.vgg_path
|
||||
BATCH_SIZE = args.batch_size
|
||||
CONTENT_WEIGHT = args.content_weight
|
||||
STYLE_WEIGHT = args.style_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
|
||||
|
||||
# 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()
|
||||
|
||||
|
||||
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
|
||||
#ttoi_t = transforms.Compose([
|
||||
# transforms.Normalize([-103.939, -116.779, -123.68],[1,1,1])])
|
||||
|
||||
# Remove the batch_size dimension
|
||||
tensor = tensor.squeeze()
|
||||
#img = ttoi_t(tensor)
|
||||
img = tensor.cpu().numpy()
|
||||
|
||||
# Transpose from [C, H, W] -> [H, W, C]
|
||||
img = img.transpose(1, 2, 0)
|
||||
return img
|
||||
|
||||
# VGG
|
||||
class VGG16(nn.Module):
|
||||
def __init__(self, vgg_path="models/vgg16-00b39a1b.pth"):
|
||||
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
|
||||
for param in self.features.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
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
|
||||
|
||||
# Transformer
|
||||
class TransformerNetworkClass(nn.Module):
|
||||
"""Feedforward Transformation Network without Tanh
|
||||
reference: https://arxiv.org/abs/1603.08155
|
||||
exact architecture: https://cs.stanford.edu/people/jcjohns/papers/fast-style/fast-style-supp.pdf
|
||||
"""
|
||||
def __init__(self):
|
||||
super(TransformerNetworkClass, self).__init__()
|
||||
self.ConvBlock = nn.Sequential(
|
||||
ConvLayer(3, 32, 9, 1),
|
||||
nn.ReLU(),
|
||||
ConvLayer(32, 64, 3, 2),
|
||||
nn.ReLU(),
|
||||
ConvLayer(64, 128, 3, 2),
|
||||
nn.ReLU()
|
||||
)
|
||||
self.ResidualBlock = nn.Sequential(
|
||||
ResidualLayer(128, 3),
|
||||
ResidualLayer(128, 3),
|
||||
ResidualLayer(128, 3),
|
||||
ResidualLayer(128, 3),
|
||||
ResidualLayer(128, 3)
|
||||
)
|
||||
self.DeconvBlock = nn.Sequential(
|
||||
DeconvLayer(128, 64, 3, 2, 1),
|
||||
nn.ReLU(),
|
||||
DeconvLayer(64, 32, 3, 2, 1),
|
||||
nn.ReLU(),
|
||||
ConvLayer(32, 3, 9, 1, norm="None")
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.ConvBlock(x)
|
||||
x = self.ResidualBlock(x)
|
||||
out = self.DeconvBlock(x)
|
||||
return out
|
||||
|
||||
class ConvLayer(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size, stride, norm="instance"):
|
||||
super(ConvLayer, self).__init__()
|
||||
# Padding Layers
|
||||
padding_size = kernel_size // 2
|
||||
self.reflection_pad = nn.ReflectionPad2d(padding_size)
|
||||
|
||||
# Convolution Layer
|
||||
self.conv_layer = nn.Conv2d(in_channels, out_channels, kernel_size, stride)
|
||||
|
||||
# Normalization Layers
|
||||
self.norm_type = norm
|
||||
if (norm=="instance"):
|
||||
self.norm_layer = nn.InstanceNorm2d(out_channels, affine=True)
|
||||
elif (norm=="batch"):
|
||||
self.norm_layer = nn.BatchNorm2d(out_channels, affine=True)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.reflection_pad(x)
|
||||
x = self.conv_layer(x)
|
||||
if (self.norm_type=="None"):
|
||||
out = x
|
||||
else:
|
||||
out = self.norm_layer(x)
|
||||
return out
|
||||
|
||||
class ResidualLayer(nn.Module):
|
||||
"""
|
||||
Deep Residual Learning for Image Recognition
|
||||
|
||||
https://arxiv.org/abs/1512.03385
|
||||
"""
|
||||
def __init__(self, channels=128, kernel_size=3):
|
||||
super(ResidualLayer, self).__init__()
|
||||
self.conv1 = ConvLayer(channels, channels, kernel_size, stride=1)
|
||||
self.relu = nn.ReLU()
|
||||
self.conv2 = ConvLayer(channels, channels, kernel_size, stride=1)
|
||||
|
||||
def forward(self, x):
|
||||
identity = x # preserve residual
|
||||
out = self.relu(self.conv1(x)) # 1st conv layer + activation
|
||||
out = self.conv2(out) # 2nd conv layer
|
||||
out = out + identity # add residual
|
||||
return out
|
||||
|
||||
class DeconvLayer(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size, stride, output_padding, norm="instance"):
|
||||
super(DeconvLayer, self).__init__()
|
||||
|
||||
# Transposed Convolution
|
||||
padding_size = kernel_size // 2
|
||||
self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding_size, output_padding)
|
||||
|
||||
# Normalization Layers
|
||||
self.norm_type = norm
|
||||
if (norm=="instance"):
|
||||
self.norm_layer = nn.InstanceNorm2d(out_channels, affine=True)
|
||||
elif (norm=="batch"):
|
||||
self.norm_layer = nn.BatchNorm2d(out_channels, affine=True)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv_transpose(x)
|
||||
if (self.norm_type=="None"):
|
||||
out = x
|
||||
else:
|
||||
out = self.norm_layer(x)
|
||||
return out
|
||||
|
||||
|
||||
|
||||
def train():
|
||||
# Seeds
|
||||
torch.manual_seed(SEED)
|
||||
torch.cuda.manual_seed(SEED)
|
||||
np.random.seed(SEED)
|
||||
random.seed(SEED)
|
||||
|
||||
# Device
|
||||
device = ("cuda" if torch.cuda.is_available() else "cpu")
|
||||
print(f"Train.py: Device is {device}")
|
||||
|
||||
# Dataset and Dataloader
|
||||
transform = transforms.Compose([
|
||||
transforms.Resize(TRAIN_IMAGE_SIZE),
|
||||
transforms.CenterCrop(TRAIN_IMAGE_SIZE),
|
||||
transforms.ToTensor(),
|
||||
transforms.Lambda(lambda x: x.mul(255))
|
||||
])
|
||||
train_dataset = datasets.ImageFolder(DATASET_PATH, transform=transform)
|
||||
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
|
||||
print(f"Train.py: train_loader length: {len(train_loader)}")
|
||||
|
||||
# Load networks
|
||||
TransformerNetwork = TransformerNetworkClass().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 = style_tensor.add(imagenet_neg_mean)
|
||||
B, C, H, W = style_tensor.shape
|
||||
style_features = VGG(style_tensor.expand([BATCH_SIZE, C, H, W]))
|
||||
style_gram = {}
|
||||
for key, value in style_features.items():
|
||||
style_gram[key] = gram(value)
|
||||
|
||||
# Optimizer settings
|
||||
optimizer = optim.AdamW(TransformerNetwork.parameters(), lr=ADAM_LR, fused=True, amsgrad=True)
|
||||
|
||||
# Loss trackers
|
||||
content_loss_history = []
|
||||
style_loss_history = []
|
||||
total_loss_history = []
|
||||
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
|
||||
optimizer.zero_grad()
|
||||
|
||||
# Generate images and get features
|
||||
content_batch = content_batch[:,[2,1,0]].to(device)
|
||||
generated_batch = TransformerNetwork(content_batch)
|
||||
content_features = VGG(content_batch.add(imagenet_neg_mean))
|
||||
generated_features = VGG(generated_batch.add(imagenet_neg_mean))
|
||||
|
||||
# Content Loss
|
||||
MSELoss = nn.MSELoss().to(device)
|
||||
content_loss = CONTENT_WEIGHT * MSELoss(generated_features['relu2_2'], content_features['relu2_2'])
|
||||
batch_content_loss_sum += content_loss
|
||||
|
||||
# Style Loss
|
||||
style_loss = 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
|
||||
batch_style_loss_sum += style_loss.item()
|
||||
|
||||
# Total Loss
|
||||
total_loss = content_loss + style_loss
|
||||
batch_total_loss_sum += total_loss.item()
|
||||
|
||||
# Backprop and Weight Update
|
||||
total_loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
# Save Model and Print Losses
|
||||
if (((batch_count-1)%SAVE_MODEL_EVERY == 0) or (batch_count==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("Time elapsed:\t{} seconds".format(time.time()-start_time))
|
||||
|
||||
# Save Model
|
||||
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"
|
||||
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)
|
||||
|
||||
# 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
|
||||
print("Done Training the Transformer Network!")
|
||||
print("Training Time: {} seconds".format(stop_time-start_time))
|
||||
print("========Content Loss========")
|
||||
print(content_loss_history)
|
||||
print("========Style Loss========")
|
||||
print(style_loss_history)
|
||||
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()
|
||||
Reference in New Issue
Block a user