Update train.py

train.py

@@ -4,6 +4,7 @@ import torch.optim as optim
 import numpy as np
 import albumentations as albu
 import argparse
+import datetime
 
 from utils.utils import open_json, weights_init, weights_init_spectr, generate_mask
 from model.models import Colorizer, Generator, Content, Discriminator
@@ -26,8 +27,8 @@ def parse_args():
 def get_transforms():
     return albu.Compose([albu.RandomCrop(512, 512, always_apply = True), albu.HorizontalFlip(p = 0.5)], p = 1.)
 
-def get_dataloaders(data_path, transforms, batch_size, fine_tuning):
-    train_dataset = TrainDataset(data_path, transforms)
+def get_dataloaders(data_path, transforms, batch_size, fine_tuning, mult_number):
+    train_dataset = TrainDataset(data_path, transforms, mult_number)
     train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size = batch_size, shuffle = True)
     
     if fine_tuning:
@@ -58,12 +59,17 @@ def set_weights(colorizer, discriminator):
     
     discriminator.apply(weights_init_spectr)
     
-def get_losses():
-    L1_loss = nn.L1Loss()
-    mse_loss = nn.MSELoss()
-    bce_loss = nn.BCEWithLogitsLoss()
+def generator_loss(disc_output, true_labels, main_output, guide_output, real_image, content_gen, content_true, dist_loss = nn.L1Loss(), content_dist_loss = nn.MSELoss(), class_loss = nn.BCEWithLogitsLoss()):    
+    sim_loss_full = dist_loss(main_output, real_image)
+    sim_loss_guide = dist_loss(guide_output, real_image)
     
-    return L1_loss, bce_loss, mse_loss
+    adv_loss = class_loss(disc_output, true_labels)
+    
+    content_loss = content_dist_loss(content_gen, content_true)
+    
+    sum_loss = 10 * (sim_loss_full + 0.9 * sim_loss_guide)  + adv_loss + content_loss
+    
+    return sum_loss
     
 def get_optimizers(colorizer, discriminator, generator_lr, discriminator_lr):
     optimizerG = optim.Adam(colorizer.generator.parameters(), lr = generator_lr, betas=(0.5, 0.9))
@@ -71,8 +77,196 @@ def get_optimizers(colorizer, discriminator, generator_lr, discriminator_lr):
     
     return optimizerG, optimizerD
 
+def generator_step(inputs, colorizer, discriminator, content,  loss_function, optimizer, device, white_penalty = True):
+    for p in discriminator.parameters():
+        p.requires_grad = False  
+    for p in colorizer.generator.parameters():
+        p.requires_grad = True    
+
+    colorizer.generator.zero_grad()
+
+    bw, color, hint, dfm = inputs  
+    bw, color, hint, dfm = bw.to(device), color.to(device), hint.to(device), dfm.to(device)
 
+    fake, guide = colorizer(torch.cat([bw, dfm, hint], 1))
 
+    logits_fake = discriminator(fake)
+    y_real = torch.ones((bw.size(0), 1), device = device)
+
+    content_fake = content(fake)
+    with torch.no_grad():
+        content_true = content(color)
+
+    generator_loss = loss_function(logits_fake, y_real, fake, guide, color, content_fake, content_true)
+    
+    if white_penalty:
+        mask = (~((color > 0.85).float().sum(dim = 1) == 3).unsqueeze(1).repeat((1, 3, 1, 1 ))).float()
+        white_zones = mask * (fake + 1) / 2
+        white_penalty = (torch.pow(white_zones.sum(dim = 1), 2).sum(dim = (1, 2)) / (mask.sum(dim = (1, 2, 3)) + 1)).mean()
+        
+        generator_loss += white_penalty
+
+    generator_loss.backward()
+
+    optimizer.step()
+    
+    return generator_loss.item()
+
+def discriminator_step(inputs, colorizer, discriminator, optimizer, device, loss_function = nn.BCEWithLogitsLoss()):
+    
+    for p in discriminator.parameters():
+        p.requires_grad = True 
+    for p in colorizer.generator.parameters():
+        p.requires_grad = False
+
+    discriminator.zero_grad()
+
+    bw, color, hint, dfm = inputs  
+    bw, color, hint, dfm = bw.to(device), color.to(device), hint.to(device), dfm.to(device)
+    
+    y_real = torch.full((bw.size(0), 1), 0.9, device = device)
+
+    y_fake = torch.zeros((bw.size(0), 1), device = device)
+
+    with torch.no_grad():
+        fake_color, _ = colorizer(torch.cat([bw, dfm, hint], 1))
+        fake_color.detach()
+
+    logits_fake = discriminator(fake_color)
+    logits_real = discriminator(color)
+
+    fake_loss = loss_function(logits_fake, y_fake)
+    real_loss = loss_function(logits_real, y_real)
+    
+    discriminator_loss = real_loss + fake_loss
+
+    discriminator_loss.backward()
+    optimizer.step()
+    
+    return discriminator_loss.item()
+
+def decrease_lr(optimizer, rate):
+    for group in optimizer.param_groups:
+        group['lr'] /= rate 
+        
+def set_lr(optimizer, value):
+    for group in optimizer.param_groups:
+        group['lr'] = value
+        
+def train(colorizer, discriminator, content, dataloader, epochs, colorizer_optimizer, discriminator_optimizer, lr_decay_epoch = -1, device = 'cpu'):
+    colorizer.generator.train()
+    discriminator.train()
+
+    disc_step  = True
+    
+    for epoch in range(epochs):
+        if (epoch == lr_decay_epoch):
+            decrease_lr(colorizer_optimizer, 10)
+            decrease_lr(discriminator_optimizer, 10)
+            
+        sum_disc_loss = 0
+        sum_gen_loss = 0
+        
+        for n, inputs in enumerate(dataloader):
+            if n % 5 == 0:
+                print(datetime.datetime.now().time())
+                print('Step : %d Discr loss: %.4f Gen loss : %.4f \n'%(n, sum_disc_loss / (n // 2 + 1), sum_gen_loss / (n // 2 + 1)))
+            
+            
+            if disc_step:
+                step_loss = discriminator_step(inputs, colorizer, discriminator, discriminator_optimizer, device)
+                sum_disc_loss += step_loss
+            else:
+                step_loss = generator_step(inputs, colorizer, discriminator, content, generator_loss, colorizer_optimizer, device)
+                sum_gen_loss += step_loss
+                
+            disc_step = disc_step ^ True
+    
+    
+        print(datetime.datetime.now().time())
+        print('Epoch : %d Discr loss: %.4f Gen loss : %.4f \n'%(epoch, sum_disc_loss / (n // 2 + 1), sum_gen_loss / (n // 2 + 1)))
+        
+        
+def fine_tuning_step(data_iter, colorizer, discriminator, gen_optimizer, disc_optimizer, device, loss_function = nn.BCEWithLogitsLoss()):
+    
+    for p in discriminator.parameters():
+        p.requires_grad = True 
+    for p in colorizer.generator.parameters():
+        p.requires_grad = False
+        
+    for cur_disc_step in range(5):
+        discriminator.zero_grad()
+        
+        bw, dfm, color_for_real = data_iter.next()
+        bw, dfm, color_for_real = bw.to(device), dfm.to(device), color_for_real.to(device)
+        
+        y_real = torch.full((bw.size(0), 1), 0.9, device = device)
+        y_fake = torch.zeros((bw.size(0), 1), device = device)
+
+        empty_hint = torch.zeros(bw.shape[0], 4, bw.shape[2] , bw.shape[3] ).float().to(device)
+        
+        with torch.no_grad():
+            fake_color_manga, _ = colorizer(torch.cat([bw, dfm, empty_hint ], 1))
+            fake_color_manga.detach()
+            
+        logits_fake = discriminator(fake_color_manga)
+        logits_real = discriminator(color_for_real)
+
+        fake_loss = loss_function(logits_fake, y_fake)
+        real_loss = loss_function(logits_real, y_real)
+        discriminator_loss = real_loss + fake_loss
+
+        discriminator_loss.backward()
+        disc_optimizer.step()
+        
+        
+    for p in discriminator.parameters():
+        p.requires_grad = False  
+    for p in colorizer.generator.parameters():
+        p.requires_grad = True
+        
+    colorizer.generator.zero_grad()    
+
+    bw, dfm, _ = data_iter.next()
+    bw, dfm = bw.to(device), dfm.to(device)
+    
+    y_real = torch.ones((bw.size(0), 1), device = device)
+    
+    empty_hint = torch.zeros(bw.shape[0], 4, bw.shape[2] , bw.shape[3]).float().to(device)
+    
+    fake_manga, _ = colorizer(torch.cat([bw, dfm, empty_hint], 1))
+
+    logits_fake = discriminator(fake_manga)
+    adv_loss = loss_function(logits_fake, y_real)
+    
+    generator_loss = adv_loss
+    
+    generator_loss.backward()
+    gen_optimizer.step()
+    
+    
+        
+def fine_tuning(colorizer, discriminator, content, dataloader, iterations, colorizer_optimizer, discriminator_optimizer, data_iter, device = 'cpu'):
+    colorizer.generator.train()
+    discriminator.train()
+    
+    disc_step = True
+    
+    for n, inputs in enumerate(dataloader):
+        
+        if n == iterations:
+            return
+        
+        if disc_step:
+            discriminator_step(inputs, colorizer, discriminator, discriminator_optimizer, device)
+        else:
+            generator_step(inputs, colorizer, discriminator, content, generator_loss, colorizer_optimizer, device)
+
+        disc_step = disc_step ^ True
+        
+        if n % 10 == 5:
+            fine_tuning_step(data_iter, colorizer, discriminator, colorizer_optimizer, discriminator_optimizer, device)
+    
 if __name__ == '__main__':
     args = parse_args()
     config = open_json('configs/train_config.json')
@@ -84,13 +278,17 @@ if __name__ == '__main__':
         
     augmentations = get_transforms()
     
-    train_dataloader, ft_dataloader = get_dataloaders(args.path, augmentations, config['batch_size'], args.fine_tuning)
+    train_dataloader, ft_dataloader = get_dataloaders(args.path, augmentations, config['batch_size'], args.fine_tuning, config['number_of_mults'])
     
     colorizer, discriminator, content = get_models(device)
     set_weights(colorizer, discriminator)
     
-    l1_loss, bce_loss, mse_loss = get_losses()
-    
     gen_optimizer, disc_optimizer = get_optimizers(colorizer, discriminator, config['generator_lr'], config['discriminator_lr'])
     
-    
+    train(colorizer, discriminator, content, train_dataloader, config['epochs'], gen_optimizer, disc_optimizer, config['lr_decrease_epoch'], device)
+    
+    if args.fine_tuning:
+        set_lr(gen_optimizer, config["finetuning_generator_lr"])
+        fine_tuning(colorizer, discriminator, content, train_dataloader, config['finetuning_iterations'], gen_optimizer, disc_optimizer, iter(ft_dataloader), device)
+    
+    torch.save(colorizer.generator.state_dict(), str(datetime.datetime.now().time()))