Complete all modules without debug

Author: lbytony

base/__init__.py

@@ -0,0 +1,3 @@
+from .base_data_loader import BaseDataLoader
+from .base_model import BaseModel
+from .base_trainer import BaseTrainer

base/base_data_loader.py

@@ -0,0 +1,106 @@
+from copy import copy
+
+import numpy as np
+
+
+class BaseDataLoader:
+    """
+    Base class for all data loaders
+    """
+
+    def __init__(self, config):
+        self.config = config
+        self.batch_size = config['data_loader']['batch_size']
+        self.shuffle = config['data_loader']['shuffle']
+        self.num_workers = config['data_loader']['workers']
+        self.activate = config['data_loader']['activate']
+        self.val_rate = config['split_rate']['validation']
+        self.test_rate = config['split_rate']['test']
+        self.batch_idx = 0
+
+    def __iter__(self):
+        """
+        :return: Iterator
+        """
+        assert self.__len__() > 0
+        self.batch_idx = 0
+        if self.shuffle:
+            self._shuffle_data()
+        return self
+
+    def __next__(self):
+        """
+        :return: Next batch
+        """
+        packed = self._pack_data()
+        if self.batch_idx < self.__len__():
+            batch = packed[self.batch_idx * self.batch_size:(self.batch_idx + 1) * self.batch_size]
+            self.batch_idx = self.batch_idx + 1
+            return self._unpack_data(batch)
+        else:
+            raise StopIteration
+
+    def __len__(self):
+        """
+        :return: Total number of batches
+        """
+        return self._n_samples() // self.batch_size
+
+    def _n_samples(self):
+        """
+        :return: Total number of samples
+        """
+        return NotImplementedError
+
+    def _pack_data(self):
+        """
+        Pack all data into a list/tuple/ndarray/...
+
+        :return: Packed data in the data loader
+        """
+        return NotImplementedError
+
+    def _unpack_data(self, packed):
+        """
+        Unpack packed data (from _pack_data())
+
+        :param packed: Packed data
+        :return: Unpacked data
+        """
+        return NotImplementedError
+
+    def _update_data(self, unpacked):
+        """
+        Update data member in the data loader
+
+        :param unpacked: Unpacked data (from _update_data())
+        """
+        return NotImplementedError
+
+    def _shuffle_data(self):
+        """
+        Shuffle data members in the data loader
+        """
+        packed = self._pack_data()
+        rand_idx = np.random.permutation(len(packed))
+        packed = [packed[i] for i in rand_idx]
+        self._update_data(self._unpack_data(packed))
+
+    def split_validation(self):
+        """
+        Split validation data from data loader based on self.config['validation']
+        """
+        validation_split = self.config['validation']['validation_split']
+        shuffle = self.config['validation']['shuffle']
+        if validation_split == 0.0:
+            return None
+        if shuffle:
+            self._shuffle_data()
+        valid_data_loader = copy(self)
+        split = int(self._n_samples() * validation_split)
+        packed = self._pack_data()
+        train_data = self._unpack_data(packed[split:])
+        val_data = self._unpack_data(packed[:split])
+        valid_data_loader._update_data(val_data)
+        self._update_data(train_data)
+        return valid_data_loader

base/base_trainer.py

@@ -0,0 +1,185 @@
+import json
+import logging
+import math
+import os
+
+import tensorflow as tf
+from tensorboardX import SummaryWriter
+from tensorflow import keras
+
+from utils.utils import ensure_dir
+
+
+class BaseTrainer:
+    """
+    Base class for all trainers
+    """
+
+    def __init__(self, model, loss, metrics, resume, config, train_logger=None):
+        self.config = config
+        self.logger = logging.getLogger(self.__class__.__name__)
+        self.model = model
+        self.loss = loss
+        self.metrics = metrics
+        self.name = config['name']
+        self.epochs = config['trainer']['epochs']
+        self.save_freq = config['trainer']['save_freq']
+        self.verbosity = config['trainer']['verbosity']
+        self.summaryWriter = SummaryWriter()
+
+        if tf.test.is_gpu_available():
+            if config['cuda']:
+                self.with_cuda = True
+                self.gpus = {i: item for i, item in enumerate(self.config['gpus'])}
+                device = 'cuda'
+            else:
+                self.with_cuda = False
+                device = 'cpu'
+        else:
+            self.logger.warning('Warning: There\'s no CUDA support on this machine, training is performed on CPU.')
+            self.with_cuda = False
+            device = 'cpu'
+
+        self.device = tf.device(device)
+        self.model.to(self.device)
+
+        self.logger.debug('Model is initialized.')
+        self._log_memory_usage()
+
+        self.train_logger = train_logger
+
+        self.optimizer = self.model.optimize(config['optimizer_type'], config['optimizer'])
+
+        self.lr_scheduler = getattr(
+            keras.callbacks.LearningRateScheduler,
+            config['lr_scheduler_type'], None)
+        if self.lr_scheduler:
+            self.lr_scheduler = self.lr_scheduler(self.optimizer, **config['lr_scheduler'])
+            self.lr_scheduler_freq = config['lr_scheduler_freq']
+        self.monitor = config['trainer']['monitor']
+        self.monitor_mode = config['trainer']['monitor_mode']
+        assert self.monitor_mode == 'min' or self.monitor_mode == 'max'
+        self.monitor_best = math.inf if self.monitor_mode == 'min' else -math.inf
+        self.start_epoch = 1
+        self.checkpoint_dir = os.path.join(config['trainer']['save_dir'], self.name)
+        ensure_dir(self.checkpoint_dir)
+        json.dump(config, open(os.path.join(self.checkpoint_dir, 'config.json'), 'w'),
+                  indent=4, sort_keys=False)
+        if resume:
+            self._resume_checkpoint(resume)
+
+    def train(self):
+        """
+        Full training logic
+        """
+        print(self.epochs)
+        for epoch in range(self.start_epoch, self.epochs + 1):
+            try:
+                result = self._train_epoch(epoch)
+            except tf.errors.UnavailableError:
+                self._log_memory_usage()
+
+            log = {'epoch': epoch}
+            for key, value in result.items():
+                if key == 'metrics':
+                    for i, metric in enumerate(self.metrics):
+                        log[metric.__name__] = result['metrics'][i]
+                elif key == 'val_metrics':
+                    for i, metric in enumerate(self.metrics):
+                        log['val_' + metric.__name__] = result['val_metrics'][i]
+                else:
+                    log[key] = value
+            if self.train_logger is not None:
+                self.train_logger.add_entry(log)
+                if self.verbosity >= 1:
+                    for key, value in log.items():
+                        self.logger.info('    {:15s}: {}'.format(str(key), value))
+            if (self.monitor_mode == 'min' and log[self.monitor] < self.monitor_best) \
+                    or (self.monitor_mode == 'max' and log[self.monitor] > self.monitor_best):
+                self.monitor_best = log[self.monitor]
+                self._save_checkpoint(epoch, log, save_best=True)
+            if epoch % self.save_freq == 0:
+                self._save_checkpoint(epoch, log)
+            if self.lr_scheduler:
+                self.lr_scheduler.step()
+                lr = self.lr_scheduler.get_lr()[0]
+                self.logger.info('New Learning Rate: {:.8f}'.format(lr))
+
+            self.summaryWriter.add_scalars('Train', {'train_' + self.monitor: result[self.monitor],
+                                                     'val_' + self.monitor: result[self.monitor]}, epoch)
+        self.summaryWriter.close()
+
+    # TODO Not Available
+    def _log_memory_usage(self):
+        if not self.with_cuda:
+            return
+
+        template = """Memory Usage: \n{}"""
+        usage = []
+        for deviceID, device in self.gpus.items():
+            deviceID = int(deviceID)
+            # allocated = torch.cuda.memory_allocated(deviceID) / (1024 * 1024)
+            # cached = torch.cuda.memory_cached(deviceID) / (1024 * 1024)
+
+            # usage.append('    CUDA: {}  Allocated: {} MB Cached: {} MB \n'.format(device, allocated, cached))
+
+        content = ''.join(usage)
+        content = template.format(content)
+
+        self.logger.debug(content)
+
+    def _train_epoch(self, epoch):
+        """
+        Training logic for an epoch
+
+        :param epoch: Current epoch number
+        """
+        raise NotImplementedError
+
+    def _save_checkpoint(self, epoch, log, save_best=False):
+        """
+        Saving checkpoints
+
+        :param epoch: current epoch number
+        :param log: logging information of the epoch
+        :param save_best: if True, rename the saved checkpoint to 'model_best.pth.tar'
+        """
+        arch = type(self.model).__name__
+        state = {
+            'arch': arch,
+            'epoch': epoch,
+            'logger': self.train_logger,
+            'state_dict': self.model.state_dict(),
+            'optimizer': self.optimizer.state_dict(),
+            'monitor_best': self.monitor_best,
+            'config': self.config
+        }
+        filename = os.path.join(self.checkpoint_dir, 'checkpoint-epoch{:03d}-loss-{:.4f}.pth.tar'
+                                .format(epoch, log['loss']))
+        tf.saved_model.save(state, filename)
+        if save_best:
+            os.rename(filename, os.path.join(self.checkpoint_dir, 'model_best.pth.tar'))
+            self.logger.info("Saving current best: {} ...".format('model_best.pth.tar'))
+        else:
+            self.logger.info("Saving checkpoint: {} ...".format(filename))
+
+    def _resume_checkpoint(self, resume_path):
+        """
+        Resume from saved checkpoints
+
+        :param resume_path: Checkpoint path to be resumed
+        """
+        self.logger.info("Loading checkpoint: {} ...".format(resume_path))
+        checkpoint = tf.saved_model.load(resume_path)
+        self.start_epoch = checkpoint['epoch'] + 1
+        self.monitor_best = checkpoint['monitor_best']
+        self.model.load_state_dict(checkpoint['state_dict'])
+        self.optimizer.load_state_dict(checkpoint['optimizer'])
+        if self.with_cuda:
+            for state in self.optimizer.state.values():
+                for k, v in state.items():
+                    if isinstance(v, tf.Tensor):
+                        state[k] = v.cuda(tf.device('cuda'))
+        self.train_logger = checkpoint['logger']
+        # self.config = checkpoint['config']
+        self.logger.info("Checkpoint '{}' (epoch {}) loaded".format(resume_path, self.start_epoch))

config.json

@@ -4,16 +4,23 @@
   "gpus": [
     0
   ],
+  "training": "True",
   "data_loader": {
-    "dataset": "ICDAR 2019 - LSVT",
-    "data_dir": "F:\\Code\\HealthHelper\\Dataset\\ICDAR 2019 - LSVT",
+    "activate": 0,
+    "datasets": [
+      {
+        "name": "ICDAR 2019 LSVT",
+        "data_dir": "F:\\Code\\HealthHelper\\Dataset\\ICDAR 2019 - LSVT",
+        "have_test": "False"
+      }
+    ],
     "batch_size": 128,
     "shuffle": true,
     "workers": 0
   },
-  "validation": {
-    "validation_split": 0.2,
-    "shuffle": true
+  "split_rate": {
+    "validation": 0.1,
+    "test": 0.1
   },
   "lr_scheduler_type": "",
   "lr_scheduler_freq": 10000,

data_loader/data_loader.py

@@ -0,0 +1,51 @@
+import tensorflow as tf
+from sklearn import model_selection
+
+from base import BaseDataLoader
+from .datasets import ICDAR2019Dataset
+
+
+class ICDAR2019DataLoaderFactory(BaseDataLoader):
+
+    def __init__(self, config):
+        super(ICDAR2019DataLoaderFactory, self).__init__(config)
+        dataRoot = config['data_loader']['datasets'][self.activate]['data_dir']
+        self.workers = config['data_loader']['workers']
+        self.have_test = config['data_loader']['datasets'][self.activate]['have_test']
+        dataset = ICDAR2019Dataset(dataRoot)
+        self.allDataset = dataset.loadData()
+
+        if self.have_test:
+            self.trainDataset, self.testDataset = self.train_val_split(self.allDataset)
+            self.trainDataset, self.valDataset = self.train_val_split(self.trainDataset)
+        else:
+            self.trainDataset, self.valDataset = self.train_val_split(self.allDataset)
+
+    def train(self):
+        trainLoader = tf.data.Dataset.from_tensor_slices(self.trainDataset)
+        # trainLoader = torchdata.DataLoader(self.trainDataset, num_workers=self.num_workers,
+        #                                    batch_size=self.batch_size,
+        #                                    shuffle=self.shuffle, collate_fn=collate_fn)
+        return trainLoader
+
+    def val(self):
+        # valLoader = torchdata.DataLoader(self.valDataset, num_workers=self.num_workers, batch_size=self.batch_size,
+        #                                  shuffle=shuffle, collate_fn=collate_fn)
+        valLoader = tf.data.Dataset.from_tensor_slices(self.valDataset)
+        return valLoader
+
+    def train_val_split(self, dataset):
+        """
+
+        :param dataset: dataset
+        :return:
+        """
+        train, val = model_selection.train_test_split(dataset[0], tuple(dataset[1:]), test_size=self.val_rate)
+        return train, val
+
+    def train_test_split(self, dataset):
+        train, test = model_selection.train_test_split(dataset[0], dataset[1:], test_size=self.test_rate)
+        return train, test
+
+    def split_validation(self):
+        raise NotImplementedError