clean code, add load config from .yaml file

Author: KhaLee2307

Dockerfile

@@ -28,10 +28,10 @@ ENV PATH /root/miniconda3/envs/strda/bin:$PATH
 
 # install dependencies
 RUN pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113
-RUN pip install opencv-python==4.4.0.46 Pillow==7.2.0 opencv-python-headless==4.5.1.48 lmdb tqdm nltk
+RUN pip install opencv-python==4.4.0.46 Pillow==7.2.0 opencv-python-headless==4.5.1.48 lmdb tqdm nltk six pyyaml
 
 RUN apt-get update
 RUN apt-get install -y ffmpeg libsm6 libxext6
 
 # get repository
-WORKDIR /home
+WORKDIR /home

config/default.yaml

@@ -0,0 +1,34 @@
+# Data Processing
+batch_max_length: 25 # maximum-label-length
+imgH: 32 # the height of the input image
+imgW: 100 # the width of the input image
+character: "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~" # character label
+
+# Model Architecture
+num_fiducial: 20 # number of fiducial points of TPS-STN"
+input_channel: 3 # the number of input channel of Feature extractor
+output_channel: 512 # the number of output channel of Feature extractor
+hidden_size: 256 # the size of the LSTM hidden state
+
+# Optimizer
+lr: 0.001 # learning rate, 0.001 for Adam
+weight_decay: 0.01 # weight decay, 0.01 for Adam
+
+# Experiment
+manual_seed: 111 # for random seed setting
+
+# Training
+grad_clip: 5 # gradient clipping value
+workers: 4 # number of data loading workers
+
+# HDGE
+decay_epoch: 100 # epoch from which to start lr decay
+load_height: 48
+load_width: 160
+crop_height: 32
+crop_width: 100
+lamda: 10
+idt_coef: 0.5
+ngf: 64 # of gen filters in first conv layer
+ndf: 64 # of discrim filters in first conv layer
+norm: "instance" # instance normalization or batch normalization

requirements.txt

@@ -1,5 +1,7 @@
+six
 lmdb
 tqdm
 nltk
+pyyaml
 pillow
-opencv-python
+opencv-python

source/HDGE.py

@@ -1,5 +1,6 @@
 import os
 import itertools
+
 import numpy as np
 from tqdm import tqdm
 
@@ -8,11 +9,11 @@
 from torch.autograd import Variable
 from torch.utils.data import Subset
 
-import utils.utils_HDGE as utils
-
 from .ops import set_grad
 from .dataset import AlignCollateHDGE, hierarchical_dataset
 
+import utils.utils_HDGE as utils
+
 from modules.generators import define_Gen
 from modules.discriminators import define_Dis
 
@@ -47,7 +48,7 @@ def __init__(self,args):
             os.makedirs(args.checkpoint_dir)
 
         try:
-            ckpt = utils.load_checkpoint('%s/latest.ckpt' % (args.checkpoint_dir))
+            ckpt = utils.load_checkpoint('%s/HDGE_gen_dis.ckpt' % (args.checkpoint_dir))
             self.start_epoch = ckpt['epoch']
             self.Da.load_state_dict(ckpt['Da'])
             self.Db.load_state_dict(ckpt['Db'])
@@ -73,7 +74,7 @@ def train(self,args):
             source_data,
             batch_size=args.batch_size,
             shuffle=True,
-            num_workers=args.num_workers,
+            num_workers=args.workers,
             collate_fn=myAlignCollate,
             pin_memory=False,
             drop_last=True,
@@ -82,7 +83,7 @@ def train(self,args):
             target_data_adjust,
             batch_size=args.batch_size,
             shuffle=True,
-            num_workers=args.num_workers,
+            num_workers=args.workers,
             collate_fn=myAlignCollate,
             pin_memory=False,
             drop_last=True,
@@ -182,8 +183,8 @@ def train(self,args):
                 b_dis_loss.backward()
                 self.d_optimizer.step()
 
-            print("\nEpoch: (%3d/%3d) | Gen Loss: %0.4f | Dis Loss: %0.4f" % 
-                    (epoch, args.epochs, gen_loss,a_dis_loss+b_dis_loss))
+            print("\nEpoch: (%3d/%3d) | Gen Loss: %0.4f | Dis Loss: %0.4f\n" % 
+                    (epoch + 1, args.epochs, gen_loss,a_dis_loss+b_dis_loss))
 
             # override the latest checkpoint
             utils.save_checkpoint({'epoch': epoch + 1,

source/dataset.py

@@ -1,8 +1,9 @@
 import os
 import sys
 import six
-import PIL
 import lmdb
+
+import PIL
 from PIL import Image
 
 import torch
@@ -15,13 +16,13 @@
 _STD_IMAGENET  = torch.tensor([0.229, 0.224, 0.225])
 
 
-def get_dataloader(opt, dataset, batch_size, shuffle = False, mode = "label"):
+def get_dataloader(args, dataset, batch_size, shuffle = False, mode = "label"):
     """
     Get dataloader for each dataset
 
     Parameters
     ----------
-    opt: argparse.ArgumentParser().parse_args()
+    args: argparse.ArgumentParser().parse_args()
     dataset: torch.utils.data.Dataset
     batch_size: int
     shuffle: boolean
@@ -32,23 +33,23 @@ def get_dataloader(opt, dataset, batch_size, shuffle = False, mode = "label"):
     """
 
     if mode == "raw":
-        myAlignCollate = AlignCollateRaw(opt)
+        myAlignCollate = AlignCollateRaw(args)
     else:
-        myAlignCollate = AlignCollate(opt, mode)
+        myAlignCollate = AlignCollate(args, mode)
 
     data_loader = DataLoader(
             dataset,
             batch_size=batch_size,
             shuffle=shuffle,
-            num_workers=opt.workers,
+            num_workers=args.workers,
             collate_fn=myAlignCollate,
             pin_memory=False,
             drop_last=False,
         )
     return data_loader
 
 
-def hierarchical_dataset(root, opt, mode="label", drop_data=[]):
+def hierarchical_dataset(root, args, mode="label", drop_data=[]):
     """ select_data='/' contains all sub-directory of root directory """
     dataset_list = []
     dataset_log = f"dataset_root:    {root}\t dataset:"
@@ -72,10 +73,10 @@ def hierarchical_dataset(root, opt, mode="label", drop_data=[]):
     for dirpath in listdir:
         if mode == "raw":
             # load data without label
-            dataset = LmdbDataset_raw(dirpath, opt)
+            dataset = LmdbDataset_raw(dirpath, args)
         else:
             # load data with label
-            dataset = LmdbDataset(dirpath, opt)
+            dataset = LmdbDataset(dirpath, args)
         sub_dataset_log = f"sub-directory:\t/{os.path.relpath(dirpath, root)}\t num samples: {len(dataset)}"
         print(sub_dataset_log)
         dataset_log += f"{sub_dataset_log}\n"
@@ -113,15 +114,15 @@ def __getitem__(self, index):
 
 class AlignCollate(object):
     """ Transform data to the same format """
-    def __init__(self, opt, mode = "label"):
-        self.opt = opt
+    def __init__(self, args, mode = "label"):
+        self.args = args
         # resize image
         if (mode == "adapt" or mode == "supervised"):
             self.transform = Rand_augment()
         else:
             self.transform = torchvision.transforms.Compose([])
 
-        self.resize = ResizeNormalize(opt)
+        self.resize = ResizeNormalize(args)
         print("Use Text_augment", self.transform)
 
     def __call__(self, batch):
@@ -135,10 +136,10 @@ def __call__(self, batch):
 
 class AlignCollateRaw(object):
     """ Transform data to the same format """
-    def __init__(self, opt):
-        self.opt = opt
+    def __init__(self, args):
+        self.args = args
         # resize image
-        self.transform = ResizeNormalize(opt)
+        self.transform = ResizeNormalize(args)
 
     def __call__(self, batch):
         images = batch
@@ -151,20 +152,20 @@ def __call__(self, batch):
 
 class AlignCollateHDGE(object):
     """ Transform data to the same format """
-    def __init__(self, opt, infer=False):
-        self.opt = opt
+    def __init__(self, args, infer=False):
+        self.args = args
 
         # for transforming the input image
         if infer == False:
             transform = torchvision.transforms.Compose(
                 [torchvision.transforms.RandomHorizontalFlip(),
-                torchvision.transforms.Resize((opt.load_height,opt.load_width)),
-                torchvision.transforms.RandomCrop((opt.crop_height,opt.crop_width)),
+                torchvision.transforms.Resize((args.load_height,args.load_width)),
+                torchvision.transforms.RandomCrop((args.crop_height,args.crop_width)),
                 torchvision.transforms.ToTensor(),
                 torchvision.transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
         else:
             transform = torchvision.transforms.Compose(
-                [torchvision.transforms.Resize((opt.crop_height,opt.crop_width)),
+                [torchvision.transforms.Resize((args.crop_height,args.crop_width)),
                 torchvision.transforms.ToTensor(),
                 torchvision.transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
 
@@ -181,10 +182,10 @@ def __call__(self, batch):
 
 class LmdbDataset(Dataset):
     """ Load data from Lmdb file with label """
-    def __init__(self, root, opt):
+    def __init__(self, root, args):
 
         self.root = root
-        self.opt = opt
+        self.args = args
         self.env = lmdb.open(
             root,
             max_readers=32,
@@ -207,7 +208,7 @@ def __init__(self, root, opt):
 
                 # length filtering
                 length_of_label = len(label)
-                if length_of_label > opt.batch_max_length:
+                if length_of_label > args.batch_max_length:
                     continue
 
                 self.filtered_index_list.append(index)
@@ -236,18 +237,18 @@ def __getitem__(self, index):
             except IOError:
                 print(f"Corrupted image for {index}")
                 # make dummy image and dummy label for corrupted image.
-                img = PIL.Image.new("RGB", (self.opt.imgW, self.opt.imgH))
+                img = PIL.Image.new("RGB", (self.args.imgW, self.args.imgH))
                 label = "[dummy_label]"
 
         return (img, label)
 
 
 class LmdbDataset_raw(Dataset):
     """ Load data from Lmdb file without label """
-    def __init__(self, root, opt):
+    def __init__(self, root, args):
 
         self.root = root
-        self.opt = opt
+        self.args = args
         self.env = lmdb.open(
             root,
             max_readers=32,
@@ -284,27 +285,21 @@ def __getitem__(self, index):
             except IOError:
                 print(f"Corrupted image for {img_key}")
                 # make dummy image for corrupted image.
-                img = PIL.Image.new("RGB", (self.opt.imgW, self.opt.imgH))
+                img = PIL.Image.new("RGB", (self.args.imgW, self.args.imgH))
 
         return img
 
 
 class ResizeNormalize(object):
 
-    def __init__(self, opt):
-        self.opt = opt
+    def __init__(self, args):
+        self.args = args
         _transforms = []
 
-        _transforms.append(
-            torchvision.transforms.Resize((self.opt.imgH, self.opt.imgW),
-                               interpolation=torchvision.transforms.InterpolationMode.BICUBIC))
+        _transforms.append(torchvision.transforms.Resize((self.args.imgH, self.args.imgW),
+                            interpolation=torchvision.transforms.InterpolationMode.BICUBIC))
         _transforms.append(torchvision.transforms.ToTensor())
-        if self.opt.use_IMAGENET_norm:
-            _transforms.append(torchvision.transforms.Normalize(mean=_MEAN_IMAGENET,
-                                                    std=_STD_IMAGENET))
-        else:
-            _transforms.append(torchvision.transforms.Normalize(mean=[0.5, 0.5, 0.5],
-                                                                std=[0.5, 0.5, 0.5]))
+        _transforms.append(torchvision.transforms.Normalize(mean=_MEAN_IMAGENET, std=_STD_IMAGENET))
         self._transforms = torchvision.transforms.Compose(_transforms)
 
     def __call__(self, image):