# Loading Dataset of celebrity face images
#dataset = ImageFolder('C:\\Users\\Shawn\\Desktop\\NYU\\LTP_SVG\\one_image', transform=img_transform) # laptop
dataset = ImageFolder('/home/so1463/LearningToPaint/baseline/one_image', transform=img_transform) # cluster
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
# AutoEncoder model
encoder = AutoEncoder().to(device)
#if os.path.exists('./AutoEncoder.pth'):
# model.load_state_dict(torch.load('/home/so1463/LearningToPaint/baseline/AutoEncoder.pth'))
# Define our RNN model
RNN = RNN().to(device)
# 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(encoder.parameters(), lr=learning_rate)
criterion = nn.MSELoss(reduction='sum')
loss_plot = []
###############################################################################
#################################
# Training ######################
#writer = TensorBoard("../train_log/")
import torch.optim as optim
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--debug', metavar='fn', default="", help="Dump outputs into file")
parser.add_argument('--script', default=False, help="Script the model")
args = parser.parse_args()
torch.manual_seed(1337)
np.random.seed(1337)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
criterion = nn.MSELoss()
net = FCN()
if args.script:
net = torch.jit.script(net)
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")
if use_cuda:
net.cuda()
from DRL.wgan import *
from utils.util import hard_update, soft_update, to_numpy
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# CoordConv implementation: https://arxiv.org/abs/1807.03247
coord = torch.zeros([1, 2, 128, 128])
for i in range(128):
for j in range(128):
coord[0, 0, i, j] = i / 127.0
coord[0, 1, i, j] = j / 127.0
coord = coord.to(device)
criterion = nn.MSELoss()
Decoder = FCN()
Decoder.load_state_dict(torch.load('../renderer.pkl'))
def decode(x, canvas, width=128):
# x: (B * 5) * (10 + 3), 5: 5 steps, 10+3: action output
x = x.view(-1, action_dim + 3)
stroke = 1 - Decoder(x[:, :action_dim])
stroke = stroke.view(-1, width, width, 1)
color_stroke = stroke * x[:, -3:].view(-1, 1, 1, 3)
stroke = stroke.permute(0, 3, 1, 2)
color_stroke = color_stroke.permute(0, 3, 1, 2)
stroke = stroke.view(-1, n_frames_per_step, 1, width, width)
color_stroke = color_stroke.view(-1, n_frames_per_step, 3, width, width)
for i in range(n_frames_per_step):
canvas = canvas * (1 - stroke[:, i]) + color_stroke[:, i]
return canvas