import cv2 import torch import numpy as np import torch.nn as nn import torch.nn.functional as F from utils.tensorboard import TensorBoard from Renderer.model import FCN from Renderer.stroke_gen import * writer = TensorBoard("../train_log/191104_renderer/") import torch.optim as optim criterion = nn.MSELoss() net = FCN() optimizer = optim.Adam(net.parameters(), lr=3e-6) batch_size = 64 use_cuda = torch.cuda.is_available() step = 0 def save_model(): if use_cuda: net.cpu() # torch.save(net.state_dict(), "../renderer.pkl") torch.save(net.state_dict(), "model/renderer.pkl") if use_cuda: net.cuda()
transform=img_transform) # cluster dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True) # AutoEncoder model model = AutoEncoder().to(device) #if os.path.exists('./AutoEncoder.pth'): # model.load_state_dict(torch.load('/home/so1463/LearningToPaint/baseline/AutoEncoder.pth')) ############################### # or use the with no_grad block in our loop when we call the save_image function # Freeze weights of the renderer renderer = FCN().to(device) renderer.load_state_dict(torch.load(args.renderer)) renderer = renderer.to(device).eval() for p in renderer.parameters(): p.requires_grad = True # Define optimizer and loss function optimizer = optim.Adam(model.parameters(), lr=learning_rate) criterion = nn.MSELoss(reduction='sum') loss_plot = [] ############################################################################### ################################# # Training ###################### ################################# imgid = 1 for epoch in range(num_epochs): # each training epoch