change model definitions and training

Author: thegenerativegeneration

config/vox-256-finetune.yaml

@@ -0,0 +1,78 @@
+dataset_params:
+  root_dir: ./video-preprocessing/vox2-768
+  frame_shape: 256,256,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: 40
+  num_repeats: 10
+  epoch_milestones: [15, 30]
+  lr_generator: 2.0e-4
+  batch_size: 16
+  scales: [1, 0.5, 0.25, 0.125]
+  dataloader_workers: 12
+  checkpoint_freq: 50
+  dropout_epoch: 2
+  dropout_maxp: 0.3
+  dropout_startp: 0.1
+  dropout_inc_epoch: 10
+  bg_start: 5
+  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: 100
+  num_repeats: 1
+  batch_size: 8
+  dataloader_workers: 6
+  checkpoint_freq: 1
+  epoch_milestones: [10, 20]
+  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.yaml

@@ -56,6 +56,10 @@ train_params:
     equivariance_value: 10
     warp_loss: 10
     bg: 10
+  optimizer: 'adam'
+  optimizer_params:
+    betas: [ 0.5, 0.999 ]
+    weight_decay: 1e-4
 
 train_avd_params:
   num_epochs: 100

config/vox-512-finetune.yaml

@@ -1,7 +1,7 @@
 # Use this file to finetune from a pretrained 256x256 model
 dataset_params:
-  root_dir: vox
-  frame_shape: null
+  root_dir: ./video-preprocessing/vox2-768
+  frame_shape: 512,512,3
   id_sampling: True
   augmentation_params:
     flip_param:
@@ -35,20 +35,20 @@ model_params:
 
 
 train_params:
-  num_epochs: 100
-  num_repeats: 10
-  epoch_milestones: [70, 90]
+  num_epochs: 30
+  num_repeats: 4
+  epoch_milestones: [20]
   # Higher LR seems to bring problems when finetuning
   lr_generator: 2.0e-5
   batch_size: 4
   scales: [1, 0.5, 0.25, 0.125]
   dataloader_workers: 6
-  checkpoint_freq: 2
-  dropout_epoch: 0
+  checkpoint_freq: 5
+  dropout_epoch: 2
   dropout_maxp: 0.3
   dropout_startp: 0.1
-  dropout_inc_epoch: 10
-  bg_start: 0
+  dropout_inc_epoch: 1
+  bg_start: 5
   transform_params:
     sigma_affine: 0.05
     sigma_tps: 0.005
@@ -58,13 +58,17 @@ train_params:
     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: 4
   dataloader_workers: 6
-  checkpoint_freq: 2
+  checkpoint_freq: 10
   epoch_milestones: [10, 20]
   lr: 1.0e-3
   lambda_shift: 1

config/vox-768-finetune.yaml

@@ -1,6 +1,6 @@
 # Use this file to finetune from a pretrained 256x256 model
 dataset_params:
-  root_dir: vox_768
+  root_dir: ./video-preprocessing/vox2-768
   frame_shape: null
   id_sampling: True
   augmentation_params:
@@ -35,20 +35,20 @@ model_params:
 
 
 train_params:
-  num_epochs: 100
-  num_repeats: 1
-  epoch_milestones: [70, 90]
+  visualize_model: False
+  num_epochs: 40
+  num_repeats: 4
   # Higher LR seems to bring problems when finetuning
   lr_generator: 2.0e-5
-  batch_size: 1
+  batch_size: 2
   scales: [1, 0.5, 0.25, 0.125]
-  dataloader_workers: 6
-  checkpoint_freq: 1
+  dataloader_workers: 8
+  checkpoint_freq: 2
   dropout_epoch: 0
   dropout_maxp: 0.3
   dropout_startp: 0.1
   dropout_inc_epoch: 10
-  bg_start: 0
+  bg_start: 5
   transform_params:
     sigma_affine: 0.05
     sigma_tps: 0.005
@@ -58,6 +58,10 @@ train_params:
     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
@@ -73,4 +77,4 @@ train_avd_params:
 visualizer_params:
   kp_size: 5
   draw_border: True
-  colormap: 'gist_rainbow'
+  colormap: 'gist_rainbow'

frames_dataset.py

@@ -67,6 +67,8 @@ def __init__(self, root_dir, frame_shape=(256, 256, 3), id_sampling=False, is_tr
                  random_seed=0, pairs_list=None, augmentation_params=None):
         self.root_dir = root_dir
         self.videos = os.listdir(root_dir)
+        if type(frame_shape) == str:
+            frame_shape = tuple(map(int, frame_shape.split(',')))
         self.frame_shape = frame_shape
         print(self.frame_shape)
         self.pairs_list = pairs_list
@@ -115,7 +117,13 @@ def __getitem__(self, idx):
 
                 frames = os.listdir(path)
                 num_frames = len(frames)
-                frame_idx = np.sort(np.random.choice(num_frames, replace=True, size=2))
+                # use more frames that are different from each other to speed up training
+                min_frames_apart = num_frames // 4
+                first_frame_idx = np.random.choice(num_frames - min_frames_apart)
+                second_frame_idx = np.random.choice(range(first_frame_idx + min_frames_apart, num_frames))
+                frame_idx = np.array([first_frame_idx, second_frame_idx])
+                np.random.shuffle(frame_idx)
+                #frame_idx = np.sort(np.random.choice(num_frames, replace=True, size=2))
 
                 if self.frame_shape is not None:
                     resize_fn = partial(resize, output_shape=self.frame_shape)

logger.py

@@ -17,12 +17,13 @@ class Logger:
     def __init__(self, log_dir, checkpoint_freq=50, visualizer_params=None,
                  zfill_num=8, log_file_name='log.txt', models=()):
 
-        self.models = None
+        self.models = models
         self.loss_list = []
         self.cpk_dir = log_dir
         self.visualizations_dir = os.path.join(log_dir, 'train-vis')
         if not os.path.exists(self.visualizations_dir):
             os.makedirs(self.visualizations_dir)
+        print("Visualizations will be saved in %s" % self.visualizations_dir)
         self.log_file = open(os.path.join(log_dir, log_file_name), 'a')
         self.zfill_num = zfill_num
         self.visualizer = Visualizer(**visualizer_params)
@@ -46,9 +47,10 @@ def log_scores(self, loss_names):
 
     def visualize_rec(self, inp, out):
         image = self.visualizer.visualize(inp['driving'], inp['source'], out)
+        wandb.log({"image": [wandb.Image(image)]})
         imageio.imsave(os.path.join(self.visualizations_dir, "%s-rec.png" % str(self.epoch).zfill(self.zfill_num)),
                        image)
-        wandb.log({"image": [wandb.Image(image)]})
+
 
     def save_cpk(self, emergent=False):
         cpk = {k: v.state_dict() for k, v in self.models.items()}

modules/bg_motion_predictor.py

@@ -11,15 +11,24 @@ class BGMotionPredictor(nn.Module):
     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)),
+            ])
         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
         self.bg_encoder.fc = nn.Linear(num_features, 6)
         self.bg_encoder.fc.weight.data.zero_()
         self.bg_encoder.fc.bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))
 
     def forward(self, source_image, driving_image):
+
+
         bs = source_image.shape[0]
         out = torch.eye(3).unsqueeze(0).repeat(bs, 1, 1).type(source_image.type())
+
+        source_image = self.preprocess(source_image)
+        driving_image = self.preprocess(driving_image)
+
         prediction = self.bg_encoder(torch.cat([source_image, driving_image], dim=1))
         out[:, :2, :] = prediction.view(bs, 2, 3)
         return out

modules/keypoint_detector.py

@@ -15,9 +15,13 @@ def __init__(self, num_tps, **kwargs):
         self.fg_encoder = models.resnet18(weights=torchvision.models.ResNet18_Weights.DEFAULT)
         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)),
+        ])
 
 
     def forward(self, image):
+        image = self.preprocess(image)
 
         fg_kp = self.fg_encoder(image)
         bs, _, = fg_kp.shape

modules/util.py

@@ -150,10 +150,10 @@ def __init__(self, in_features, kernel_size, padding):
 
     def forward(self, x):
         out = self.norm1(x)
-        out = F.relu(out)
+        out = F.mish(out)
         out = self.conv1(out)
         out = self.norm2(out)
-        out = F.relu(out)
+        out = F.mish(out)
         out = self.conv2(out)
         out += x
         return out
@@ -172,10 +172,10 @@ def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1
         self.norm = nn.InstanceNorm2d(out_features, affine=True)
 
     def forward(self, x):
-        out = F.interpolate(x, scale_factor=2)
+        out = F.interpolate(x, scale_factor=2, mode='nearest')
         out = self.conv(out)
         out = self.norm(out)
-        out = F.relu(out)
+        out = F.mish(out)
         return out
 
 
@@ -194,7 +194,7 @@ def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1
     def forward(self, x):
         out = self.conv(x)
         out = self.norm(out)
-        out = F.relu(out)
+        out = F.mish(out)
         out = self.pool(out)
         return out
 
@@ -213,7 +213,7 @@ def __init__(self, in_features, out_features, groups=1, kernel_size=3, padding=1
     def forward(self, x):
         out = self.conv(x)
         out = self.norm(out)
-        out = F.relu(out)
+        out = F.mish(out)
         return out
 
 

run.py

@@ -18,6 +18,7 @@
 from train_avd import train_avd
 from reconstruction import reconstruction
 import os
+from torchinfo import summary
 import bitsandbytes as bnb
 
 optimizer_choices = {
@@ -37,7 +38,7 @@
     parser.add_argument("--mode", default="train", choices=["train", "reconstruction", "train_avd"])
     parser.add_argument("--log_dir", default='log', help="path to log into")
     parser.add_argument("--checkpoint", default=None, help="path to checkpoint to restore")
-    parser.add_argument("--optimizer_class", default="adam", choices=optimizer_choices.keys())
+    parser.add_argument("--detect_anomaly", action="store_true", help="detect anomaly in autograd")
 
 
     opt = parser.parse_args()
@@ -50,6 +51,9 @@
         log_dir = os.path.join(opt.log_dir, os.path.basename(opt.config).split('.')[0])
         log_dir += ' ' + strftime("%d_%m_%y_%H.%M.%S", gmtime())
 
+    if opt.detect_anomaly:
+        torch.autograd.set_detect_anomaly(True)
+
     inpainting = InpaintingNetwork(**config['model_params']['generator_params'],
                                         **config['model_params']['common_params'])
 
@@ -76,7 +80,17 @@
     if not os.path.exists(os.path.join(log_dir, os.path.basename(opt.config))):
         copy(opt.config, log_dir)
 
-    optimizer_class = optimizer_choices[opt.optimizer_class]
+    optimizer_class = optimizer_choices[config['train_params']['optimizer']]
+
+    print("Inpainting Network:")
+    summary(inpainting)
+    print("Keypoint Detector:")
+    summary(kp_detector)
+    print("Dense Motion Network:")
+    summary(dense_motion_network)
+    if bg_predictor is not None:
+        print("Background Predictor:")
+        summary(bg_predictor)
 
     if opt.mode == 'train':
         print("Training...")
@@ -90,3 +104,4 @@
         print("Reconstruction...")
         #TODO: update to accelerate
         reconstruction(config, inpainting, kp_detector, bg_predictor, dense_motion_network, opt.checkpoint, log_dir, dataset)
+