def closure(): # correct the values of updated input image input_img.data.clamp_(0, 1) optimizer.zero_grad() model(input_img) style_score = 0 content_score = 0 for sl in style_losses: style_score += sl.loss for cl in content_losses: content_score += cl.loss style_score *= style_weight content_score *= content_weight loss = style_score + content_score loss.backward() run[0] += 1 if run[0] % 5 == 0: print("run {}:".format(run)) print('Style Loss : {:4f} Content Loss: {:4f}'.format( style_score.item(), content_score.item())) print() if prev: img_handler.imshow(input_img, unloader, title='Run {}'.format(run)) if runs: run_save = os.path.join(run_path, f'run{run}.jpg') img_handler.imsave(run_save, unloader, input_img) return style_score + content_score
output_path = os.path.join(output_path, out_img) print(f'Output path: {output_path}') # add the original input image to the figure: if plt_prvs: plt.figure() img_handler.imshow(input_img, unloader, title='Input Image') # --- RUNNING THE ALGORITHM --- style_model = StyleModel(device, cnn, cnn_normalization_mean, cnn_normalization_std, style_img, content_img, style_layers=style_layers_default, content_layers=content_layers_default) output = style_model.run_style_transfer(input_img=input_img, unloader=unloader, prev=plt_prvs, runs=run_saves) # --- SHOW AND SAVE OUTPUT --- if plt_prvs: img_handler.imshow(output, unloader, title='Output Image') plt.ioff() plt.show() img_handler.imsave(output_path, unloader, output) input("Press enter to continue...")