def webcam(style_transform_path, width=1280, height=720):
"""
Captures and saves an image, perform style transfer, and again saves the styled image.
Reads the styled image and show in window.
Saving and loading SHOULD BE eliminated, however this produces too much whitening in
the "generated styled image". This may be caused by the async nature of VideoCapture,
and I don't know how to fix it.
"""
# Device
device = ("cuda" if torch.cuda.is_available() else "cpu")
# Load Transformer Network
print("Loading Transformer Network")
net = transformer.TransformerNetwork()
net.load_state_dict(torch.load(style_transform_path))
net = net.to(device)
print("Done Loading Transformer Network")
# Set webcam settings
cam = cv2.VideoCapture(0)
cam.set(3, width)
cam.set(4, height)
# Main loop
with torch.no_grad():
count = 1
while True:
# Get webcam input
ret_val, img = cam.read()
# Mirror
img = cv2.flip(img, 1)
utils.saveimg(img, str(count) + ".png")
# Free-up unneeded cuda memory
torch.cuda.empty_cache()
# Generate image
content_image = utils.load_image(str(count) + ".png")
content_tensor = utils.itot(content_image).to(device)
generated_tensor = net(content_tensor)
generated_image = utils.ttoi(generated_tensor.detach())
if (PRESERVE_COLOR):
generated_image = utils.transfer_color(content_image,
generated_image)
utils.saveimg(generated_image, str(count + 1) + ".png")
img2 = cv2.imread(str(count + 1) + ".png")
count += 2
# Show webcam
cv2.imshow('Demo webcam', img2)
if cv2.waitKey(1) == 27:
break # esc to quit
# Free-up memories
cam.release()
cv2.destroyAllWindows()
def stylize_folder(style_path, folder_containing_the_content_folder, save_folder, batch_size=1):
"""Stylizes images in a folder by batch
If the images are of different dimensions, use transform.resize() or use a batch size of 1
IMPORTANT: Put content_folder inside another folder folder_containing_the_content_folder
folder_containing_the_content_folder
content_folder
pic1.ext
pic2.ext
pic3.ext
...
and saves as the styled images in save_folder as follow:
save_folder
pic1.ext
pic2.ext
pic3.ext
...
"""
# Device
device = ("cuda" if torch.cuda.is_available() else "cpu")
# Image loader
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(lambda x: x.mul(255))
])
image_dataset = utils.ImageFolderWithPaths(folder_containing_the_content_folder, transform=transform)
image_loader = torch.utils.data.DataLoader(image_dataset, batch_size=batch_size)
# Load Transformer Network
net = transformer.TransformerNetwork()
net.load_state_dict(torch.load(style_path))
net = net.to(device)
# Stylize batches of images
with torch.no_grad():
for content_batch, _, path in image_loader:
# Free-up unneeded cuda memory
torch.cuda.empty_cache()
# Generate image
generated_tensor = net(content_batch.to(device)).detach()
# Save images
for i in range(len(path)):
generated_image = utils.ttoi(generated_tensor[i])
if (PRESERVE_COLOR):
generated_image = utils.transfer_color(content_image, generated_image)
image_name = os.path.basename(path[i])
utils.saveimg(generated_image, save_folder + image_name)
#stylize()
def webcam(style_transform_path, width=1280, height=720):
"""
Captures and saves an image, perform style transfer, and again saves the styled image.
Reads the styled image and show in window.
"""
# Device
device = "cuda" if torch.cuda.is_available() else "cpu"
# Load Transformer Network
print("Loading Transformer Network")
net = transformer.TransformerNetwork()
net.load_state_dict(torch.load(style_transform_path))
net = net.to(device)
print("Done Loading Transformer Network")
# Set webcam settings
cam = cv2.VideoCapture(0)
cam.set(3, width)
cam.set(4, height)
# Main loop
with torch.no_grad():
while True:
# Get webcam input
ret_val, img = cam.read()
# Mirror
img = cv2.flip(img, 1)
# Free-up unneeded cuda memory
torch.cuda.empty_cache()
# Generate image
content_tensor = utils.itot(img).to(device)
generated_tensor = net(content_tensor)
generated_image = utils.ttoi(generated_tensor.detach())
if PRESERVE_COLOR:
generated_image = utils.transfer_color(img, generated_image)
generated_image = generated_image / 255
# Show webcam
cv2.imshow("Demo webcam", generated_image)
if cv2.waitKey(1) == 27:
break # esc to quit
# Free-up memories
cam.release()
cv2.destroyAllWindows()
def webcam(style_transform_path, width=1280, height=720):
# Device
device = ("cuda" if torch.cuda.is_available() else "cpu")
# Load Transformer Network
print("Loading Transformer Network")
net = transformer.TransformerNetwork()
net.load_state_dict(torch.load(style_transform_path, map_location=device))
net = net.to(device)
print("Done Loading Transformer Network")
# Set webcam settings
cam = cv2.VideoCapture(0)
cam.set(3, width)
cam.set(4, height)
# Main loop
with torch.no_grad():
count = 1
while True:
# Get webcam input
ret_val, img = cam.read()
# AI Mirror
img = cv2.flip(img, 1)
# Free-up unneeded cuda memory
torch.cuda.empty_cache()
# Generate image
content_tensor = utils.itot(img).to(device)
generated_tensor = net(content_tensor)
generated_image = utils.ttoi(generated_tensor.detach())
if (PRESERVE_COLOR):
generated_image = utils.transfer_color(content_image,
generated_image)
img2 = cv2.imdecode(
cv2.imencode(".png", generated_image)[1], cv2.IMREAD_UNCHANGED)
count += 2
# Show webcam
cv2.imshow('Demo webcam', img2)
if cv2.waitKey(1) == 27:
break # esc to quit
# Free-up memories
cam.release()
cv2.destroyAllWindows()
def stylize_folder_single(style_path, content_folder, save_folder):
"""
Reads frames/pictures as follows:
content_folder
pic1.ext
pic2.ext
pic3.ext
...
and saves as the styled images in save_folder as follow:
save_folder
pic1.ext
pic2.ext
pic3.ext
...
"""
# Device
device = ("cuda" if torch.cuda.is_available() else "cpu")
# Load Transformer Network
net = transformer.TransformerNetwork()
net.load_state_dict(torch.load(style_path))
net = net.to(device)
# Stylize every frame
images = [img for img in os.listdir(content_folder) if img.endswith(".jpg")]
with torch.no_grad():
for image_name in images:
# Free-up unneeded cuda memory
torch.cuda.empty_cache()
# Load content image
content_image = utils.load_image(content_folder + image_name)
content_tensor = utils.itot(content_image).to(device)
# Generate image
generated_tensor = net(content_tensor)
generated_image = utils.ttoi(generated_tensor.detach())
if (PRESERVE_COLOR):
generated_image = utils.transfer_color(content_image, generated_image)
# Save image
utils.saveimg(generated_image, save_folder + image_name)
def stylize():
# Device
device = ("cuda" if torch.cuda.is_available() else "cpu")
# Load Transformer Network
net = transformer.TransformerNetwork()
net.load_state_dict(torch.load(STYLE_TRANSFORM_PATH))
net = net.to(device)
with torch.no_grad():
while(1):
torch.cuda.empty_cache()
print("Stylize Image~ Press Ctrl+C and Enter to close the program")
content_image_path = input("Enter the image path: ")
content_image = utils.load_image(content_image_path)
starttime = time.time()
content_tensor = utils.itot(content_image).to(device)
generated_tensor = net(content_tensor)
generated_image = utils.ttoi(generated_tensor.detach())
if (PRESERVE_COLOR):
generated_image = utils.transfer_color(content_image, generated_image)
print("Transfer Time: {}".format(time.time() - starttime))
utils.show(generated_image)
utils.saveimg(generated_image, "helloworld.jpg")
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"
# 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)
# Load networks
TransformerNetwork = transformer.TransformerNetwork().to(device)
VGG = vgg.VGG16().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))
style_image = utils.load_image(STYLE_IMAGE_PATH)
style_tensor = utils.itot(style_image).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]))
style_gram = {}
for key, value in style_features.items():
style_gram[key] = utils.gram(value)
# Optimizer settings
optimizer = optim.Adam(TransformerNetwork.parameters(), lr=ADAM_LR)
# 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(utils.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 = utils.ttoi(sample_tensor.clone().detach())
sample_image_path = (SAVE_IMAGE_PATH + "sample0_" +
str(batch_count - 1) + ".png")
utils.saveimg(sample_image, sample_image_path)
print("Saved sample tranformed image at {}".format(
sample_image_path))
# Save loss histories
content_loss_history.append(batch_total_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
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:
utils.plot_loss_hist(content_loss_history, style_loss_history,
total_loss_history)
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")
# Dataset and Dataloader
transform = transforms.Compose([
transforms.Resize(TRAIN_IMAGE_SIZE),
transforms.CenterCrop(TRAIN_IMAGE_SIZE),
transforms.ToTensor()
])
x_trainloader, x_testloader = utils.prepare_loader(DATASET_PATH,
X_CLASS,
transform=transform,
batch_size=BATCH_SIZE,
shuffle=True)
y_trainloader, y_testloader = utils.prepare_loader(DATASET_PATH,
Y_CLASS,
transform=transform,
batch_size=BATCH_SIZE,
shuffle=True)
# Load Networks
Gxy = models.Generator(64).to(device)
Gyx = models.Generator(64).to(device)
Dxy = models.Discriminator(64).to(device)
Dyx = models.Discriminator(64).to(device)
# Optimizer Settings
G_param = list(Gxy.parameters()) + list(Gyx.parameters())
G_optim = optim.Adam(G_param, lr=LR, betas=[BETA_1, BETA_2])
Dxy_optim = optim.Adam(Dxy.parameters(), lr=LR, betas=[BETA_1, BETA_2])
Dyx_optim = optim.Adam(Dyx.parameters(), lr=LR, betas=[BETA_1, BETA_2])
# Losses
from losses import real_loss, fake_loss, cycle_loss
# Fixed test samples
x_testiter = iter(x_testloader)
y_testiter = iter(y_testloader)
fixed_X = next(x_testiter)[0]
fixed_X = fixed_X.to(device)
fixed_Y = next(y_testiter)[0]
fixed_Y = fixed_Y.to(device)
# Tensor to Image
fixed_X_image = utils.ttoi(fixed_X)
fixed_Y_image = utils.ttoi(fixed_Y)
# Number of batches
x_trainiter = iter(x_trainloader)
y_trainiter = iter(y_trainloader)
iter_per_epoch = min(len(x_trainiter), len(y_trainiter))
# Training the CycleGAN
for epoch in range(1, NUM_EPOCHS + 1):
print("========Epoch {}/{}========".format(epoch, NUM_EPOCHS))
for _ in range(iter_per_epoch -
1): # -1 in case of imbalanced sizes of the last batch
# Fetch the dataset
x_real = next(x_trainiter)[0]
x_real = x_real.to(device)
y_real = next(y_trainiter)[0]
y_real = y_real.to(device)
# ========= Discriminator ==========
# In training the discriminators, we fix the generators' parameters.
# It is alright to train both discriminators seperately beceause
# their forward pass don't share any parameters with each other
# Discriminator X -> Y Adversarial Loss
Dxy_optim.zero_grad() # Zero-out gradients
Dxy_real_out = Dxy(y_real) # Dxy Forward Pass
Dxy_real_loss = real_loss(Dxy_real_out) # Dxy Eeal loss
Dxy_fake_out = Dxy(Gxy(x_real)) # Gxy produces fake-y images
Dxy_fake_loss = fake_loss(Dxy_fake_out) # Dxy Fake Loss
Dxy_loss = Dxy_real_loss + Dxy_fake_loss # Dxy Total Loss
Dxy_loss.backward() # Dxy Backprop
Dxy_optim.step() # Dxy Gradient Descent
# Discriminator Y-> X Adversarial Loss
Dyx_optim.zero_grad() # Zero-out gradients
Dyx_real_out = Dyx(x_real) # Dyx Forward Pass
Dyx_real_loss = real_loss(Dyx_real_out) # Dyx Eeal loss
Dyx_fake_out = Dyx(Gyx(y_real)) # Gyx produces fake-x images
Dyx_fake_loss = fake_loss(Dyx_fake_out) # Dyx Fake Loss
Dyx_loss = Dyx_real_loss + Dyx_fake_loss # Dyx Total Loss
Dyx_loss.backward() # Dyx Backprop
Dyx_optim.step() # Dyx Gradient Descent
# ============= Generator ==============
# Similar to training discriminator networks, in training
# generator networks, we fix discriminator networks.
# However, cycle consistency prohibits us
# from training generators seperately.
# Generator X -> Y Adversarial Loss
G_optim.zero_grad() # Zero-out gradients
Gxy_out = Gxy(x_real) # Gxy Forward Pass
D_Gxy_out = Dxy(Gxy_out) # Gxy -> Dxy Forward
Gxy_loss = real_loss(D_Gxy_out) # Gxy Real Loss
# Generator Y -> X Adversarial Loss
Gyx_out = Gyx(y_real) # Gyx Forward Pass
D_Gyx_out = Dyx(Gyx_out) # Gyx -> Dyx Forward
Gyx_loss = real_loss(D_Gyx_out) # Gyx Real Loss
# Cycle Consistency Loss
y_x_y = Gxy(Gyx(x_real)) # Reconstruct Y
yxy_cycle_loss = cycle_loss(
y_x_y, y_real) # Y-X-Y Cycle Reconstruction Loss
x_y_x = Gyx(Gxy(y_real)) # Reconstruct X
xyx_cycle_loss = cycle_loss(
x_y_x, x_real) # X-Y-X Cycle Reconstruction Loss
# Generator Total Loss
G_loss = Gxy_loss + Gyx_loss + CYCLE_WEIGHT * (xyx_cycle_loss +
yxy_cycle_loss)
G_loss.backward()
G_optim.step()
# Print Losses
print("Dxy Loss: {} Dyx Loss: {} Generator Loss: {}".format(
Dxy_loss.item(), Dyx_loss.item(), G_loss.item()))
# Generate Sample Fake Images
Gxy.eval()
Gyx.eval()
with torch.no_grad():
generated_y = Gyx(fixed_X)
generated_y_img = utils.ttoi(generated_y.clone().detach())
generated_x = Gxy(fixed_Y)
generated_x_img = utils.ttoi(generated_x.clone().detach())
H = W = TRAIN_IMAGE_SIZE
concat_y = utils.concatenate_images(fixed_Y_image, generated_y_img,
H, W)
concat_x = utils.concatenate_images(fixed_X_image, generated_x_img,
H, W)
utils.saveimg(concat_x, "generated_x.png")
utils.saveimg(concat_y, "generated_y.png")
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")
# Dataset and Dataloader
transform = transforms.Compose([
transforms.Resize(TRAIN_IMAGE_SIZE),
transforms.CenterCrop(TRAIN_IMAGE_SIZE),
# transforms.Grayscale(num_output_channels=3),
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)
# Load networks
TransformerNetwork = transformer.TransformerNetwork().to(device)
if USE_LATEST_CHECKPOINT is True:
files = glob.glob(
"/home/clng/github/fast-neural-style-pytorch/models/checkpoint*")
if len(files) == 0:
print("use latest checkpoint but no checkpoint found")
else:
files.sort(key=os.path.getmtime, reverse=True)
latest_checkpoint_path = files[0]
print("using latest checkpoint %s" % (latest_checkpoint_path))
params = torch.load(latest_checkpoint_path, map_location=device)
TransformerNetwork.load_state_dict(params)
VGG = vgg.VGG19().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)
style_image = utils.load_image(STYLE_IMAGE_PATH)
if ADJUST_BRIGHTNESS == "1":
style_image = cv2.cvtColor(style_image, cv2.COLOR_BGR2GRAY)
style_image = utils.hist_norm(style_image,
[0, 64, 96, 128, 160, 192, 255],
[0, 0.05, 0.15, 0.5, 0.85, 0.95, 1],
inplace=True)
elif ADJUST_BRIGHTNESS == "2":
style_image = cv2.cvtColor(style_image, cv2.COLOR_BGR2GRAY)
style_image = cv2.equalizeHist(style_image)
elif ADJUST_BRIGHTNESS == "3":
a = 1
# hsv = cv2.cvtColor(style_image, cv2.COLOR_BGR2HSV)
# hsv = utils.auto_brightness(hsv)
# style_image = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
style_image = ensure_three_channels(style_image)
sname = os.path.splitext(os.path.basename(STYLE_IMAGE_PATH))[0] + "_train"
cv2.imwrite(
"/home/clng/datasets/bytenow/neural_styles/{s}.jpg".format(s=sname),
style_image)
style_tensor = utils.itot(style_image,
max_size=TRAIN_STYLE_SIZE).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]))
style_gram = {}
for key, value in style_features.items():
style_gram[key] = utils.gram(value)
# Optimizer settings
optimizer = optim.Adam(TransformerNetwork.parameters(), lr=ADAM_LR)
# 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['relu3_4'],
content_features['relu3_4'])
batch_content_loss_sum += content_loss
# Style Loss
style_loss = 0
for key, value in generated_features.items():
s_loss = MSELoss(utils.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 = utils.ttoi(sample_tensor.clone().detach())
sample_image_path = SAVE_IMAGE_PATH + "sample0_" + str(
batch_count - 1) + ".png"
utils.saveimg(sample_image, sample_image_path)
print("Saved sample tranformed image at {}".format(
sample_image_path))
# Save loss histories
content_loss_history.append(batch_total_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
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 + STYLE_NAME + ".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):
utils.plot_loss_hist(content_loss_history, style_loss_history,
total_loss_history)