succeeded to reproduce mocogan 🎉

Author: raahii

.gitignore

@@ -3,5 +3,6 @@ __pycache__
 deploy.sh
 
 data/*
+result/*
 notebooks/*
 !.gitkeep

result/.gitkeep

@@ -14,7 +14,7 @@
 
 class VideoDataset(Dataset):
     def __init__(self, dataset_path, preprocess_func, video_length=16, image_size=64, \
-                       mode="train"):
+                        number_limit=-1, mode="train"):
         # TODO: currently, mode only support 'train'
 
         root_path = dataset_path / 'preprocessed' / mode
@@ -28,16 +28,21 @@ def __init__(self, dataset_path, preprocess_func, video_length=16, image_size=64
                 raise e
 
         # collect video folder paths
-        video_list = []
         with open(root_path/"list.txt") as f:
-            for line in f.readlines():
-                # append [color_path, depth_path, n_frames]
-                color_path, depth_path, n_frames = line.strip().split(" ")
-                video_list.append([
-                    root_path / color_path,
-                    root_path / depth_path,
-                    int(n_frames)
-                ])
+            lines = f.readlines()
+
+        if number_limit != -1:
+            lines = lines[:number_limit]
+        
+        video_list = []
+        for line in lines:
+            # append [color_path, depth_path, n_frames]
+            color_path, depth_path, n_frames = line.strip().split(" ")
+            video_list.append([
+                root_path / color_path,
+                root_path / depth_path,
+                int(n_frames)
+            ])
 
         self.dataset_path = dataset_path
         self.root_path = root_path

src/dataset.py

@@ -21,6 +21,7 @@ def prepare_dataset(configs):
             eval(f'preprocess_{configs["dataset"]["name"]}_dataset'),
             configs['video_length'], 
             configs['image_size'], 
+            configs["dataset"]['number_limit'], 
             )
 
 def main():

src/train.py

@@ -15,11 +15,11 @@ class Trainer(object):
     def __init__(self, dataloader, configs):
 
         self.batchsize = configs["batchsize"]
+        self.epoch_iters = len(dataloader)
         self.max_iteration = configs["iterations"]
         self.video_length = configs["video_length"]
 
         self.dataloader = dataloader
-        self.dataiter = iter(dataloader)
 
         self.log_dir = Path(configs["log_dir"]) / configs["experiment_name"]
         self.log_dir.mkdir(parents=True, exist_ok=True)
@@ -28,7 +28,7 @@ def __init__(self, dataloader, configs):
         self.tensorboard_dir.mkdir(parents=True, exist_ok=True)
 
         self.logger = Logger(self.log_dir, self.tensorboard_dir,\
-                             configs["log_interval"], len(dataloader))
+                             configs["log_interval"], self.epoch_iters)
 
         self.evaluation_interval  = configs["evaluation_interval"]
         self.log_samples_interval = configs["log_samples_interval"]
@@ -38,18 +38,6 @@ def __init__(self, dataloader, configs):
         self.device = self.use_cuda and torch.device('cuda') or torch.device('cpu')
         self.configs = configs
 
-    def sample_real_batch(self):
-        try:
-            batch = next(self.dataiter)
-        except StopIteration:
-            self.data_iter = iter(self.dataloader)
-            batch = next(self.dataiter) 
-
-        if self.use_cuda:
-            batch = batch.cuda()
-
-        return batch.float()
-
     def create_optimizer(self, model, lr, decay):
         return optim.Adam(
                 model.parameters(),
@@ -90,11 +78,11 @@ def train(self, gen, idis, vdis):
 
         # training loop
         logger = self.logger
+        dataiter = iter(self.dataloader)
         while True:
             #--------------------
             # phase generator
             #--------------------
-
             gen.train();  opt_gen.zero_grad()
 
             # fake batch
@@ -118,7 +106,9 @@ def train(self, gen, idis, vdis):
             vdis.train(); opt_vdis.zero_grad()
 
             # real batch
-            x_real = Variable(self.sample_real_batch())
+            x_real = next(dataiter).float()
+            x_real = x_real.cuda() if self.use_cuda else x_fake
+            x_real = Variable(x_real)
 
             y_real_i = idis(x_real[:,:,t_rand])
             y_real_v = vdis(x_real)
@@ -143,6 +133,9 @@ def train(self, gen, idis, vdis):
 
             iteration = self.logger.metrics["iteration"]
 
+            if iteration % (self.epoch_iters-1) == 0:
+                dataiter = iter(self.dataloader)
+
             # snapshot models
             if iteration % configs["snapshot_interval"] == 0:
                 torch.save( gen, str(self.log_dir/'gen_{:05d}.pytorch'.format(iteration)))