change model definitions and training

Author: thegenerativegeneration

config/vox-1024-finetune.yaml

@@ -0,0 +1,82 @@
+# Use this file to finetune from a pretrained 256x256 model
+dataset_params:
+  root_dir: ./video-preprocessing/vox2-768
+  frame_shape: 1024,1024,3
+  id_sampling: True
+  augmentation_params:
+    flip_param:
+      horizontal_flip: True
+      time_flip: True
+    jitter_param:
+      brightness: 0.1
+      contrast: 0.1
+      saturation: 0.1
+      hue: 0.1
+
+
+model_params:
+  common_params:
+    num_tps: 10
+    num_channels: 3
+    bg: True
+    multi_mask: True
+  generator_params:
+    block_expansion: 64
+    max_features: 512
+    num_down_blocks: 3
+  dense_motion_params:
+    block_expansion: 64
+    max_features: 1024
+    num_blocks: 5
+    scale_factor: 0.25
+  avd_network_params:
+    id_bottle_size: 128
+    pose_bottle_size: 128
+
+
+train_params:
+  num_epochs: 5
+  num_repeats: 4
+  # Higher LR seems to bring problems when finetuning
+  lr_generator: 2.0e-5
+  batch_size: 1
+  scales: [1, 0.5, 0.25, 0.125, 0.0625, 0.03125]
+  dataloader_workers: 6
+  checkpoint_freq: 5
+  dropout_epoch: 2
+  dropout_maxp: 0.3
+  dropout_startp: 0.1
+  dropout_inc_epoch: 1
+  bg_start: 81
+  freeze_kp_detector: True
+  freeze_bg_predictor: True
+  transform_params:
+    sigma_affine: 0.05
+    sigma_tps: 0.005
+    points_tps: 5
+  loss_weights:
+    perceptual: [10, 10, 10, 10, 10]
+    equivariance_value: 10
+    warp_loss: 10
+    bg: 10
+  optimizer: 'adamw'
+  optimizer_params:
+    betas: [ 0.9, 0.999 ]
+    weight_decay: 0.1
+
+
+train_avd_params:
+  num_epochs: 200
+  num_repeats: 1
+  batch_size: 1
+  dataloader_workers: 6
+  checkpoint_freq: 1
+  epoch_milestones: [140, 180]
+  lr: 1.0e-3
+  lambda_shift: 1
+  random_scale: 0.25
+
+visualizer_params:
+  kp_size: 5
+  draw_border: True
+  colormap: 'gist_rainbow'

config/vox-256-finetune.yaml

@@ -34,19 +34,20 @@ model_params:
 
 
 train_params:
-  num_epochs: 40
+  num_epochs: 5
   num_repeats: 10
-  epoch_milestones: [15, 30]
-  lr_generator: 2.0e-4
+  lr_generator: 2.0e-5
   batch_size: 16
   scales: [1, 0.5, 0.25, 0.125]
   dataloader_workers: 12
-  checkpoint_freq: 50
-  dropout_epoch: 2
+  checkpoint_freq: 10
+  dropout_epoch: 0
   dropout_maxp: 0.3
   dropout_startp: 0.1
   dropout_inc_epoch: 10
-  bg_start: 5
+  bg_start: 6
+  freeze_kp_detector: False
+  freeze_bg_predictor: True
   transform_params:
     sigma_affine: 0.05
     sigma_tps: 0.005
@@ -61,6 +62,7 @@ train_params:
     betas: [ 0.9, 0.999 ]
     weight_decay: 0.1
 
+
 train_avd_params:
   num_epochs: 100
   num_repeats: 1

config/vox-512-finetune.yaml

@@ -35,20 +35,21 @@ model_params:
 
 
 train_params:
-  num_epochs: 30
+  num_epochs: 40
   num_repeats: 4
-  epoch_milestones: [20]
   # Higher LR seems to bring problems when finetuning
-  lr_generator: 2.0e-5
+  lr_generator: 2.0e-4
   batch_size: 4
-  scales: [1, 0.5, 0.25, 0.125]
+  scales: [1, 0.5, 0.25, 0.125, 0.0625]
   dataloader_workers: 6
   checkpoint_freq: 5
   dropout_epoch: 2
   dropout_maxp: 0.3
   dropout_startp: 0.1
   dropout_inc_epoch: 1
-  bg_start: 5
+  bg_start: 41
+  freeze_kp_detector: True
+  freeze_bg_predictor: True
   transform_params:
     sigma_affine: 0.05
     sigma_tps: 0.005

config/vox-768-finetune.yaml

@@ -1,7 +1,7 @@
 # Use this file to finetune from a pretrained 256x256 model
 dataset_params:
   root_dir: ./video-preprocessing/vox2-768
-  frame_shape: null
+  frame_shape: 768,768,3
   id_sampling: True
   augmentation_params:
     flip_param:
@@ -36,19 +36,21 @@ model_params:
 
 train_params:
   visualize_model: False
-  num_epochs: 40
-  num_repeats: 4
+  num_epochs: 80
+  num_repeats: 10
   # Higher LR seems to bring problems when finetuning
-  lr_generator: 2.0e-5
+  lr_generator: 3.0e-5
   batch_size: 2
-  scales: [1, 0.5, 0.25, 0.125]
+  scales: [1, 0.5, 0.25, 0.125, 0.0625]
   dataloader_workers: 8
   checkpoint_freq: 2
   dropout_epoch: 0
   dropout_maxp: 0.3
   dropout_startp: 0.1
   dropout_inc_epoch: 10
-  bg_start: 5
+  bg_start: 81
+  freeze_kp_detector: True
+  freeze_bg_predictor: True
   transform_params:
     sigma_affine: 0.05
     sigma_tps: 0.005

modules/bg_motion_predictor.py

@@ -12,7 +12,7 @@ def __init__(self):
         super(BGMotionPredictor, self).__init__()
         self.bg_encoder = models.resnet18(weights=torchvision.models.ResNet18_Weights.DEFAULT)
         self.preprocess = torchvision.transforms.Compose([
-            torchvision.transforms.Resize((256, 256)),
+            torchvision.transforms.Resize((256, 256), antialias=True),
             ])
         self.bg_encoder.conv1 = nn.Conv2d(6, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
         num_features = self.bg_encoder.fc.in_features

modules/keypoint_detector.py

@@ -16,7 +16,7 @@ def __init__(self, num_tps, **kwargs):
         num_features = self.fg_encoder.fc.in_features
         self.fg_encoder.fc = nn.Linear(num_features, num_tps*5*2)
         self.preprocess = torchvision.transforms.Compose([
-            torchvision.transforms.Resize((256, 256)),
+            torchvision.transforms.Resize((256, 256), antialias=True),
         ])
 
 

save_model_only.py

@@ -36,7 +36,7 @@
 
 
     bg_predictor = None
-    if (config['model_params']['common_params']['bg']):
+    if 'bg_predictor' in checkpoint:
         bg_predictor = BGMotionPredictor()
 
     avd_network = None

train.py

@@ -15,6 +15,7 @@
 
 accelerator = Accelerator()
 
+
 def train(config, inpainting_network, kp_detector, bg_predictor, dense_motion_network, checkpoint, log_dir, dataset,
           optimizer_class=torch.optim.Adam
           ):
@@ -44,23 +45,35 @@ def train(config, inpainting_network, kp_detector, bg_predictor, dense_motion_ne
     else:
         start_epoch = 0
 
-
-
+    freeze_kp_detector = train_params.get('freeze_kp_detector', False)
+    freeze_bg_predictor = train_params.get('freeze_bg_predictor', False)
+    if freeze_kp_detector:
+        print('freeze kp detector')
+        kp_detector.eval()
+        for param in kp_detector.parameters():
+            param.requires_grad = False
+    if freeze_bg_predictor:
+        print('freeze bg predictor')
+        bg_predictor.eval()
+        for param in bg_predictor.parameters():
+            param.requires_grad = False
 
     if 'num_repeats' in train_params or train_params['num_repeats'] != 1:
         dataset = DatasetRepeater(dataset, train_params['num_repeats'])
     dataloader = DataLoader(dataset, batch_size=train_params['batch_size'], shuffle=True,
                             num_workers=train_params['dataloader_workers'], drop_last=True)
 
     scheduler_optimizer = OneCycleLR(optimizer, max_lr=train_params['lr_generator'],
-                                     total_steps=(len(dataset) // train_params['batch_size']) * train_params['num_epochs'],
-                                      last_epoch=start_epoch-1)
+                                     total_steps=(len(dataset) // train_params['batch_size']) * train_params[
+                                         'num_epochs'],
+                                     last_epoch=start_epoch - 1)
 
     scheduler_bg_predictor = None
     if bg_predictor:
         scheduler_bg_predictor = OneCycleLR(optimizer_bg_predictor, max_lr=train_params['lr_generator'],
-                                     total_steps=(len(dataset) // train_params['batch_size']) * train_params['num_epochs'],
-                                      last_epoch=start_epoch-1)
+                                            total_steps=(len(dataset) // train_params['batch_size']) * train_params[
+                                                'num_epochs'],
+                                            last_epoch=start_epoch - 1)
         bg_predictor, optimizer_bg_predictor = accelerator.prepare(bg_predictor, optimizer_bg_predictor)
 
     generator_full = GeneratorFullModel(kp_detector, bg_predictor, dense_motion_network, inpainting_network,
@@ -75,16 +88,21 @@ def train(config, inpainting_network, kp_detector, bg_predictor, dense_motion_ne
     if train_params.get('visualize_model', False):
         # visualize graph
         sample = next(iter(dataloader))
-        draw_graph(generator_full, input_data=[sample, 100], save_graph=True, directory=log_dir, graph_name='generator_full')
-        draw_graph(kp_detector, input_data=[sample['driving']], save_graph=True, directory=log_dir, graph_name='kp_detector')
+        draw_graph(generator_full, input_data=[sample, 100], save_graph=True, directory=log_dir,
+                   graph_name='generator_full')
+        draw_graph(kp_detector, input_data=[sample['driving']], save_graph=True, directory=log_dir,
+                   graph_name='kp_detector')
         kp_driving = kp_detector(sample['driving'])
         kp_source = kp_detector(sample['source'])
         bg_param = bg_predictor(sample['source'], sample['driving'])
-        dense_motion_param = {'source_image': sample['source'], 'kp_driving':  kp_driving, 'kp_source': kp_source, 'bg_param': bg_param,
-                                                      'dropout_flag' : False, 'dropout_p' : 0.0}
+        dense_motion_param = {'source_image': sample['source'], 'kp_driving': kp_driving, 'kp_source': kp_source,
+                              'bg_param': bg_param,
+                              'dropout_flag': False, 'dropout_p': 0.0}
         dense_motion = dense_motion_network(**dense_motion_param)
-        draw_graph(dense_motion_network, input_data=dense_motion_param, save_graph=True, directory=log_dir, graph_name='dense_motion_network')
-        draw_graph(inpainting_network, input_data=[sample['source'], dense_motion], save_graph=True, directory=log_dir, graph_name='inpainting_network')
+        draw_graph(dense_motion_network, input_data=dense_motion_param, save_graph=True, directory=log_dir,
+                   graph_name='dense_motion_network')
+        draw_graph(inpainting_network, input_data=[sample['source'], dense_motion], save_graph=True, directory=log_dir,
+                   graph_name='inpainting_network')
 
     with Logger(log_dir=log_dir, visualizer_params=config['visualizer_params'],
                 checkpoint_freq=train_params['checkpoint_freq'],
@@ -100,14 +118,18 @@ def train(config, inpainting_network, kp_detector, bg_predictor, dense_motion_ne
 
                 clip_grad_norm_(kp_detector.parameters(), max_norm=10, norm_type=math.inf)
                 clip_grad_norm_(dense_motion_network.parameters(), max_norm=10, norm_type=math.inf)
-                if bg_predictor and epoch >= bg_start:
+                if bg_predictor and epoch >= bg_start and not freeze_bg_predictor:
                     clip_grad_norm_(bg_predictor.parameters(), max_norm=10, norm_type=math.inf)
 
                 optimizer.step()
-                optimizer.zero_grad()
-                if bg_predictor and epoch >= bg_start:
+
+                if bg_predictor and epoch >= bg_start and not freeze_bg_predictor:
                     optimizer_bg_predictor.step()
                     optimizer_bg_predictor.zero_grad()
+                    scheduler_bg_predictor.step()
+
+                optimizer.zero_grad()
+                scheduler_optimizer.step()
 
                 losses = {key: value.mean().detach().data.cpu().numpy() for key, value in losses_generator.items()}
                 lrs = {
@@ -116,23 +138,15 @@ def train(config, inpainting_network, kp_detector, bg_predictor, dense_motion_ne
                 }
                 logger.log_iter(losses=losses, others=lrs)
 
-                scheduler_optimizer.step()
-                if bg_predictor:
-                    scheduler_bg_predictor.step()
+
 
             model_save = {
-                'inpainting_network': inpainting_network,
-                'dense_motion_network': dense_motion_network,
-                'kp_detector': kp_detector,
+                'inpainting_network': accelerator.unwrap_model(inpainting_network),
+                'dense_motion_network': accelerator.unwrap_model(dense_motion_network),
+                'kp_detector': accelerator.unwrap_model(kp_detector),
                 'optimizer': optimizer,
+                'bg_predictor': accelerator.unwrap_model(bg_predictor) if bg_predictor else None,
+                'optimizer_bg_predictor': optimizer_bg_predictor
             }
-            if bg_predictor and epoch >= bg_start:
-                model_save['bg_predictor'] = bg_predictor
-                model_save['optimizer_bg_predictor'] = optimizer_bg_predictor
-
-            accelerator.save_state(log_dir)
-
 
             logger.log_epoch(epoch, model_save, inp=x, out=generated)
-
-