Hi Patrick,
Apologies for the confusion this is my entire training loop.
###Below if my Code
valid_loss_min = valid_loss_min_input
for epoch in range(start_epochs, n_epochs+1):
# initialize variables to monitor training and validation loss
train_loss = 0.0
###################
# train the model #
###################
model.train()
for batch_idx,img in enumerate(loaders['rgbimage_dataloader']):
optimizer.zero_grad()
# move to GPU
if use_cuda:
real = img[0].to(device)
sketch_image = cycle_model(real)
else:
real = img[0].to(device)
sketch_image = cycle_model(real)
###Getting output from main model#####
styleGan_output = model(real,sketch_image,Target_color.float(),noise)
###Feeding to the pretrained model########
fakeEdge_image = cycle_model(styleGan_output)
###Unormalizing Image to feed to PIL image
x_tensor = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])(styleGan_output)
min_i = x_tensor.min(dim=(1), keepdim=True).values.min(dim=(2), keepdim=True).values
max_i = x_tensor.max(dim=(1), keepdim=True).values.max(dim=(2), keepdim=True).values
x_tensor = ((x_tensor-min_i) / (max_i - min_i)) * 255
styleGAN_transformedimage = [transforms.ToPILImage()(x_) for x_ in x_tensor.type(torch.uint8)]
###Get Dominant Color
styleGAN_transformedimage = [get_dominant_color(x) for x in styleGAN_transformedimage]
Avg_color = [torch.tensor(x,dtype=float,requires_grad=False) for x in styleGAN_transformedimage]
Avg_color = [color.to(device) for color in Avg_color]
Avg_color = torch.stack(Avg_color)
Target_color_match = torch.cat(len(Avg_color)*[Target_color])
Avg_color = torch.unsqueeze(Avg_color,1)
###Loss calculation########
Edge_loss = criterion(fakeEdge_image,sketch_image)
color_loss = criterion(Avg_color,Target_color_match)
total_loss = Edge_loss + color_loss
total_loss.backward()
optimizer.step()
train_loss = train_loss + ((1 / (batch_idx + 1)) * (total_loss.data - train_loss))
# print training/validation statistics
print('Epoch: {} \tTraining Loss: {:.6f}'.format(epoch,train_loss))
del total_loss
del styleGan_output
del fakeEdge_image
del Edge_loss
torch.save(model.state_dict(),f"model_history/mu_stylegan{epoch}.pth")
This is my entire training loop
Below steps in the above code are not from the model output i.e sketch_image,Avg_color and Target_color_match.
Edge_loss = criterion(fakeEdge_image,sketch_image)
color_loss = criterion(Avg_color,Target_color_match)
Does states of all these to get saved below
torch.save(model.state_dict(),f"model_history/mu_stylegan{epoch}.pth")
Below is the way i am Evaluating:
Target_color = torch.zeros(1, 1, 3, dtype=torch.float)
Target_color[:,:,2] = 255
Target_color_blue = Target_color.to(device)
model.load_state_dict(torch.load(os.path.join(“model_history/mu_stylegan178.pth”),map_location=torch.device(‘cuda’)))
#progress_bar = tqdm(enumerate(trainA_dataloader,trainB_dataloader),total=len(trainA_dataloader))
model.eval()
outputs = []
for batch_idx, img in enumerate(loaders[‘test_rgbimage_dataloader’]):
if use_cuda:
real = img[0].to(device)
sketch_image = Cyclegan_model(real)
recon = mu_stylegan(real,sketch_image,Target_color_blue.float(),noise)
sketch_fake_test = recon.detach().cpu().numpy()
for img in sketch_fake_test:
img = img/2+0.5
plt.imshow(np.transpose(img, (1, 2, 0)))
plt.show()