Add application file

app.py

@@ -0,0 +1,35 @@
+import os
+os.system('pip install --upgrade gdown')
+import gdown
+gdown.download(id='1mYM_26qHUom_5NU7iutHneB_KHlLjL5y', output='workdir.zip')
+os.system('unzip workdir.zip')
+
+import glob
+import gradio as gr
+from demo import get_model, preprocess, postprocess, load
+from utils import Config, Logger, CharsetMapper
+
+def process_image(image):
+  config = Config('configs/train_abinet.yaml')
+  config.model_vision_checkpoint = None
+  model = get_model(config)
+  model = load(model, 'workdir/train-abinet/best-train-abinet.pth')
+  charset = CharsetMapper(filename=config.dataset_charset_path, max_length=config.dataset_max_length + 1)
+
+  img = image.convert('RGB')
+  img = preprocess(img, config.dataset_image_width, config.dataset_image_height)
+  res = model(img)
+  return postprocess(res, charset, 'alignment')[0][0]
+
+title = "Interactive demo: ABINet"
+description = "Demo for ABINet, ABINet uses a vision model and an explicit language model to recognize text in the wild, which are trained in end-to-end way. The language model (BCN) achieves bidirectional language representation in simulating cloze test, additionally utilizing iterative correction strategy. To use it, simply upload a (single-text line) image or use one of the example images below and click 'submit'. Results will show up in a few seconds."
+article = "<p style='text-align: center'><a href='https://arxiv.org/pdf/2103.06495.pdf'>Read Like Humans: Autonomous, Bidirectional and Iterative Language Modeling for Scene Text Recognition</a> | <a href='https://github.com/FangShancheng/ABINet'>Github Repo</a></p>"
+
+iface = gr.Interface(fn=process_image, 
+                     inputs=gr.inputs.Image(type="pil"), 
+                     outputs=gr.outputs.Textbox(),
+                     title=title,
+                     description=description,
+                     article=article,
+                     examples=glob.glob('figs/test/*.png'))
+iface.launch(debug=True)

callbacks.py

@@ -0,0 +1,360 @@
+import logging
+import shutil
+import time
+
+import editdistance as ed
+import torchvision.utils as vutils
+from fastai.callbacks.tensorboard import (LearnerTensorboardWriter,
+                                          SummaryWriter, TBWriteRequest,
+                                          asyncTBWriter)
+from fastai.vision import *
+from torch.nn.parallel import DistributedDataParallel
+from torchvision import transforms
+
+import dataset
+from utils import CharsetMapper, Timer, blend_mask
+
+
+class IterationCallback(LearnerTensorboardWriter):
+    "A `TrackerCallback` that monitor in each iteration."
+    def __init__(self, learn:Learner, name:str='model', checpoint_keep_num=5,
+                 show_iters:int=50, eval_iters:int=1000, save_iters:int=20000,
+                 start_iters:int=0, stats_iters=20000):
+        #if self.learn.rank is not None: time.sleep(self.learn.rank)  # keep all event files
+        super().__init__(learn, base_dir='.', name=learn.path, loss_iters=show_iters, 
+                        stats_iters=stats_iters, hist_iters=stats_iters)
+        self.name, self.bestname = Path(name).name, f'best-{Path(name).name}'
+        self.show_iters = show_iters
+        self.eval_iters = eval_iters
+        self.save_iters = save_iters
+        self.start_iters = start_iters
+        self.checpoint_keep_num = checpoint_keep_num
+        self.metrics_root = 'metrics/'  # rewrite
+        self.timer = Timer()
+        self.host = self.learn.rank is None or self.learn.rank == 0
+
+    def _write_metrics(self, iteration:int, names:List[str], last_metrics:MetricsList)->None:
+        "Writes training metrics to Tensorboard."
+        for i, name in enumerate(names):
+            if last_metrics is None or len(last_metrics) < i+1: return
+            scalar_value = last_metrics[i]
+            self._write_scalar(name=name, scalar_value=scalar_value, iteration=iteration)
+
+    def _write_sub_loss(self, iteration:int, last_losses:dict)->None:
+        "Writes sub loss to Tensorboard."
+        for name, loss in last_losses.items():
+            scalar_value = to_np(loss)
+            tag = self.metrics_root + name
+            self.tbwriter.add_scalar(tag=tag, scalar_value=scalar_value, global_step=iteration)
+
+    def _save(self, name):
+        if isinstance(self.learn.model, DistributedDataParallel):
+            tmp = self.learn.model
+            self.learn.model = self.learn.model.module
+            self.learn.save(name)
+            self.learn.model = tmp
+        else: self.learn.save(name)
+
+    def _validate(self, dl=None, callbacks=None, metrics=None, keeped_items=False):
+        "Validate on `dl` with potential `callbacks` and `metrics`."
+        dl = ifnone(dl, self.learn.data.valid_dl)
+        metrics = ifnone(metrics, self.learn.metrics)
+        cb_handler = CallbackHandler(ifnone(callbacks, []), metrics)
+        cb_handler.on_train_begin(1, None, metrics); cb_handler.on_epoch_begin()
+        if keeped_items: cb_handler.state_dict.update(dict(keeped_items=[]))
+        val_metrics = validate(self.learn.model, dl, self.loss_func, cb_handler)
+        cb_handler.on_epoch_end(val_metrics)
+        if keeped_items: return cb_handler.state_dict['keeped_items']
+        else: return cb_handler.state_dict['last_metrics']
+
+    def jump_to_epoch_iter(self, epoch:int, iteration:int)->None:
+        try:
+            self.learn.load(f'{self.name}_{epoch}_{iteration}', purge=False)
+            logging.info(f'Loaded {self.name}_{epoch}_{iteration}')
+        except: logging.info(f'Model {self.name}_{epoch}_{iteration} not found.')
+
+    def on_train_begin(self, n_epochs, **kwargs):
+        # TODO: can not write graph here
+        # super().on_train_begin(**kwargs)
+        self.best = -float('inf')
+        self.timer.tic()
+        if self.host:
+            checkpoint_path = self.learn.path/'checkpoint.yaml'
+            if checkpoint_path.exists():
+                os.remove(checkpoint_path)
+            open(checkpoint_path, 'w').close()
+        return {'skip_validate': True, 'iteration':self.start_iters}  # disable default validate
+
+    def on_batch_begin(self, **kwargs:Any)->None:
+        self.timer.toc_data()
+        super().on_batch_begin(**kwargs)
+
+    def on_batch_end(self, iteration, epoch, last_loss, smooth_loss, train, **kwargs):
+        super().on_batch_end(last_loss, iteration, train, **kwargs)
+        if iteration == 0: return
+
+        if iteration % self.loss_iters == 0:
+            last_losses = self.learn.loss_func.last_losses
+            self._write_sub_loss(iteration=iteration, last_losses=last_losses)
+            self.tbwriter.add_scalar(tag=self.metrics_root + 'lr',
+                scalar_value=self.opt.lr, global_step=iteration)
+
+        if iteration % self.show_iters == 0:
+            log_str = f'epoch {epoch} iter {iteration}: loss = {last_loss:6.4f},  ' \
+                      f'smooth loss = {smooth_loss:6.4f}'
+            logging.info(log_str)
+            # log_str = f'data time = {self.timer.data_diff:.4f}s, runing time = {self.timer.running_diff:.4f}s'
+            # logging.info(log_str)
+
+        if iteration % self.eval_iters == 0:
+            # TODO: or remove time to on_epoch_end
+            # 1. Record time 
+            log_str = f'average data time = {self.timer.average_data_time():.4f}s, ' \
+                      f'average running time = {self.timer.average_running_time():.4f}s'
+            logging.info(log_str)
+
+            # 2. Call validate
+            last_metrics = self._validate()
+            self.learn.model.train()
+            log_str = f'epoch {epoch} iter {iteration}: eval loss = {last_metrics[0]:6.4f},  ' \
+                      f'ccr = {last_metrics[1]:6.4f},  cwr = {last_metrics[2]:6.4f},  ' \
+                      f'ted = {last_metrics[3]:6.4f},  ned = {last_metrics[4]:6.4f},  ' \
+                      f'ted/w = {last_metrics[5]:6.4f}, '
+            logging.info(log_str)
+            names = ['eval_loss', 'ccr', 'cwr', 'ted', 'ned', 'ted/w']
+            self._write_metrics(iteration, names, last_metrics)
+
+            # 3. Save best model
+            current = last_metrics[2]
+            if current is not None and current > self.best:
+                logging.info(f'Better model found at epoch {epoch}, '\
+                             f'iter {iteration} with accuracy value: {current:6.4f}.')
+                self.best = current
+                self._save(f'{self.bestname}')
+
+        if iteration % self.save_iters == 0 and self.host:
+            logging.info(f'Save model {self.name}_{epoch}_{iteration}')
+            filename = f'{self.name}_{epoch}_{iteration}'
+            self._save(filename)
+
+            checkpoint_path = self.learn.path/'checkpoint.yaml'
+            if not checkpoint_path.exists():
+                open(checkpoint_path, 'w').close()
+            with open(checkpoint_path, 'r') as file:
+                checkpoints = yaml.load(file, Loader=yaml.FullLoader) or dict()
+            checkpoints['all_checkpoints'] = (
+                checkpoints.get('all_checkpoints') or list())
+            checkpoints['all_checkpoints'].insert(0, filename)
+            if len(checkpoints['all_checkpoints']) > self.checpoint_keep_num:
+                removed_checkpoint = checkpoints['all_checkpoints'].pop()
+                removed_checkpoint =  self.learn.path/self.learn.model_dir/f'{removed_checkpoint}.pth'
+                os.remove(removed_checkpoint)
+            checkpoints['current_checkpoint'] = filename
+            with open(checkpoint_path, 'w') as file:
+                yaml.dump(checkpoints, file)
+
+
+        self.timer.toc_running()
+
+    def on_train_end(self, **kwargs):
+        #self.learn.load(f'{self.bestname}', purge=False)
+        pass
+    
+    def on_epoch_end(self, last_metrics:MetricsList, iteration:int, **kwargs)->None:
+        self._write_embedding(iteration=iteration)
+
+
+class TextAccuracy(Callback):
+    _names = ['ccr', 'cwr', 'ted', 'ned', 'ted/w']
+    def __init__(self, charset_path, max_length, case_sensitive, model_eval):
+        self.charset_path = charset_path
+        self.max_length = max_length
+        self.case_sensitive = case_sensitive
+        self.charset = CharsetMapper(charset_path, self.max_length)
+        self.names = self._names
+
+        self.model_eval = model_eval or 'alignment'
+        assert self.model_eval in ['vision', 'language', 'alignment']
+ 
+    def on_epoch_begin(self, **kwargs):
+        self.total_num_char = 0.
+        self.total_num_word = 0.
+        self.correct_num_char = 0.
+        self.correct_num_word = 0.
+        self.total_ed = 0.
+        self.total_ned = 0.
+
+    def _get_output(self, last_output):
+        if isinstance(last_output, (tuple, list)): 
+            for res in last_output:
+                if res['name'] == self.model_eval: output = res
+        else: output = last_output
+        return output
+    
+    def _update_output(self, last_output, items):
+        if isinstance(last_output, (tuple, list)): 
+            for res in last_output:
+                if res['name'] == self.model_eval: res.update(items)
+        else: last_output.update(items)
+        return last_output
+
+    def on_batch_end(self, last_output, last_target, **kwargs):
+        output = self._get_output(last_output)
+        logits, pt_lengths = output['logits'], output['pt_lengths']
+        pt_text, pt_scores, pt_lengths_ = self.decode(logits)
+        assert (pt_lengths == pt_lengths_).all(), f'{pt_lengths} != {pt_lengths_} for {pt_text}'
+        last_output = self._update_output(last_output, {'pt_text':pt_text, 'pt_scores':pt_scores})
+
+        pt_text = [self.charset.trim(t) for t in pt_text]
+        label = last_target[0]
+        if label.dim() == 3: label = label.argmax(dim=-1)  # one-hot label
+        gt_text = [self.charset.get_text(l, trim=True) for l in label]
+        
+        for i in range(len(gt_text)):
+            if not self.case_sensitive:
+                gt_text[i], pt_text[i] = gt_text[i].lower(), pt_text[i].lower()
+            distance = ed.eval(gt_text[i], pt_text[i])
+            self.total_ed += distance
+            self.total_ned += float(distance) / max(len(gt_text[i]), 1)
+
+            if gt_text[i] == pt_text[i]:
+                self.correct_num_word += 1
+            self.total_num_word += 1
+            
+            for j in range(min(len(gt_text[i]), len(pt_text[i]))):
+                if gt_text[i][j] == pt_text[i][j]:
+                    self.correct_num_char += 1
+            self.total_num_char += len(gt_text[i])
+
+        return {'last_output': last_output}
+
+    def on_epoch_end(self, last_metrics, **kwargs):
+        mets = [self.correct_num_char / self.total_num_char,
+                self.correct_num_word / self.total_num_word,
+                self.total_ed,
+                self.total_ned,
+                self.total_ed / self.total_num_word]
+        return add_metrics(last_metrics, mets)
+
+    def decode(self, logit):
+        """ Greed decode """
+        # TODO: test running time and decode on GPU
+        out = F.softmax(logit, dim=2)
+        pt_text, pt_scores, pt_lengths = [], [], []
+        for o in out:
+            text = self.charset.get_text(o.argmax(dim=1), padding=False, trim=False)
+            text = text.split(self.charset.null_char)[0]  # end at end-token
+            pt_text.append(text)
+            pt_scores.append(o.max(dim=1)[0])
+            pt_lengths.append(min(len(text) + 1, self.max_length))  # one for end-token
+        pt_scores = torch.stack(pt_scores)
+        pt_lengths = pt_scores.new_tensor(pt_lengths, dtype=torch.long)
+        return pt_text, pt_scores, pt_lengths
+
+
+class TopKTextAccuracy(TextAccuracy):
+    _names = ['ccr', 'cwr']
+    def __init__(self, k, charset_path, max_length, case_sensitive, model_eval):
+        self.k = k
+        self.charset_path = charset_path
+        self.max_length = max_length
+        self.case_sensitive = case_sensitive
+        self.charset = CharsetMapper(charset_path, self.max_length)
+        self.names = self._names
+ 
+    def on_epoch_begin(self, **kwargs):
+        self.total_num_char = 0.
+        self.total_num_word = 0.
+        self.correct_num_char = 0.
+        self.correct_num_word = 0.
+
+    def on_batch_end(self, last_output, last_target, **kwargs):
+        logits, pt_lengths = last_output['logits'], last_output['pt_lengths']
+        gt_labels, gt_lengths = last_target[:]
+
+        for logit, pt_length, label, length in zip(logits, pt_lengths, gt_labels, gt_lengths):
+            word_flag = True
+            for i in range(length):
+                char_logit = logit[i].topk(self.k)[1]
+                char_label = label[i].argmax(-1)
+                if char_label in char_logit: self.correct_num_char += 1
+                else: word_flag = False
+                self.total_num_char += 1
+            if pt_length == length and word_flag:
+                self.correct_num_word += 1
+            self.total_num_word += 1
+
+    def on_epoch_end(self, last_metrics, **kwargs):
+        mets = [self.correct_num_char / self.total_num_char,
+                self.correct_num_word / self.total_num_word,
+                0., 0., 0.]
+        return add_metrics(last_metrics, mets)
+
+
+class DumpPrediction(LearnerCallback):
+
+    def __init__(self, learn, dataset, charset_path, model_eval, image_only=False, debug=False):
+        super().__init__(learn=learn)
+        self.debug = debug
+        self.model_eval = model_eval or 'alignment'
+        self.image_only = image_only
+        assert self.model_eval in ['vision', 'language', 'alignment']
+
+        self.dataset, self.root = dataset, Path(self.learn.path)/f'{dataset}-{self.model_eval}'
+        self.attn_root = self.root/'attn'
+        self.charset = CharsetMapper(charset_path)
+        if self.root.exists(): shutil.rmtree(self.root)
+        self.root.mkdir(), self.attn_root.mkdir()
+
+        self.pil = transforms.ToPILImage()
+        self.tensor = transforms.ToTensor()
+        size = self.learn.data.img_h, self.learn.data.img_w
+        self.resize = transforms.Resize(size=size, interpolation=0)
+        self.c = 0
+
+    def on_batch_end(self, last_input, last_output, last_target, **kwargs):
+        if isinstance(last_output, (tuple, list)):
+            for res in last_output:
+                if res['name'] == self.model_eval: pt_text = res['pt_text']
+                if res['name'] == 'vision': attn_scores = res['attn_scores'].detach().cpu()
+                if res['name'] == self.model_eval: logits = res['logits']
+        else:
+            pt_text = last_output['pt_text']
+            attn_scores = last_output['attn_scores'].detach().cpu()
+            logits = last_output['logits']
+
+        images = last_input[0] if isinstance(last_input, (tuple, list)) else last_input
+        images = images.detach().cpu()
+        pt_text = [self.charset.trim(t) for t in pt_text]
+        gt_label = last_target[0]
+        if gt_label.dim() == 3: gt_label = gt_label.argmax(dim=-1)  # one-hot label
+        gt_text = [self.charset.get_text(l, trim=True) for l in gt_label]
+
+        prediction, false_prediction = [], []
+        for gt, pt, image, attn, logit in zip(gt_text, pt_text, images, attn_scores, logits):
+            prediction.append(f'{gt}\t{pt}\n')
+            if gt != pt:
+                if self.debug:
+                    scores = torch.softmax(logit, dim=-1)[:max(len(pt), len(gt)) + 1]
+                    logging.info(f'{self.c} gt {gt}, pt {pt}, logit {logit.shape}, scores {scores.topk(5, dim=-1)}')
+                false_prediction.append(f'{gt}\t{pt}\n')
+
+            image = self.learn.data.denorm(image)
+            if not self.image_only:
+                image_np = np.array(self.pil(image))
+                attn_pil = [self.pil(a) for a in attn[:, None, :, :]]
+                attn = [self.tensor(self.resize(a)).repeat(3, 1, 1) for a in attn_pil]
+                attn_sum = np.array([np.array(a) for a in attn_pil[:len(pt)]]).sum(axis=0)
+                blended_sum = self.tensor(blend_mask(image_np, attn_sum))
+                blended = [self.tensor(blend_mask(image_np, np.array(a))) for a in attn_pil]
+                save_image = torch.stack([image] + attn + [blended_sum] + blended)
+                save_image = save_image.view(2, -1, *save_image.shape[1:])
+                save_image = save_image.permute(1, 0, 2, 3, 4).flatten(0, 1)
+                vutils.save_image(save_image, self.attn_root/f'{self.c}_{gt}_{pt}.jpg', 
+                                nrow=2, normalize=True, scale_each=True)
+            else:
+                self.pil(image).save(self.attn_root/f'{self.c}_{gt}_{pt}.jpg')
+            self.c += 1
+
+        with open(self.root/f'{self.model_eval}.txt', 'a') as f: f.writelines(prediction)
+        with open(self.root/f'{self.model_eval}-false.txt', 'a') as f: f.writelines(false_prediction)  

dataset.py

@@ -0,0 +1,278 @@
+import logging
+import re
+
+import cv2
+import lmdb
+import six
+from fastai.vision import *
+from torchvision import transforms
+
+from transforms import CVColorJitter, CVDeterioration, CVGeometry
+from utils import CharsetMapper, onehot
+
+
+class ImageDataset(Dataset):
+    "`ImageDataset` read data from LMDB database."
+
+    def __init__(self,
+                 path:PathOrStr,
+                 is_training:bool=True,
+                 img_h:int=32,
+                 img_w:int=100,
+                 max_length:int=25,
+                 check_length:bool=True,
+                 case_sensitive:bool=False,
+                 charset_path:str='data/charset_36.txt',
+                 convert_mode:str='RGB',
+                 data_aug:bool=True,
+                 deteriorate_ratio:float=0.,
+                 multiscales:bool=True,
+                 one_hot_y:bool=True,
+                 return_idx:bool=False,
+                 return_raw:bool=False,
+                 **kwargs):
+        self.path, self.name = Path(path), Path(path).name
+        assert self.path.is_dir() and self.path.exists(), f"{path} is not a valid directory."
+        self.convert_mode, self.check_length = convert_mode, check_length
+        self.img_h, self.img_w = img_h, img_w
+        self.max_length, self.one_hot_y = max_length, one_hot_y
+        self.return_idx, self.return_raw = return_idx, return_raw
+        self.case_sensitive, self.is_training = case_sensitive, is_training
+        self.data_aug, self.multiscales = data_aug, multiscales
+        self.charset = CharsetMapper(charset_path, max_length=max_length+1)
+        self.c = self.charset.num_classes
+
+        self.env = lmdb.open(str(path), readonly=True, lock=False, readahead=False, meminit=False)
+        assert self.env, f'Cannot open LMDB dataset from {path}.'
+        with self.env.begin(write=False) as txn:
+            self.length = int(txn.get('num-samples'.encode()))
+
+        if self.is_training and self.data_aug:
+            self.augment_tfs = transforms.Compose([
+                CVGeometry(degrees=45, translate=(0.0, 0.0), scale=(0.5, 2.), shear=(45, 15), distortion=0.5, p=0.5),
+                CVDeterioration(var=20, degrees=6, factor=4, p=0.25),
+                CVColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.1, p=0.25)
+            ])
+        self.totensor = transforms.ToTensor()
+    
+    def __len__(self): return self.length
+
+    def _next_image(self, index):
+        next_index = random.randint(0, len(self) - 1)
+        return self.get(next_index)
+
+    def _check_image(self, x, pixels=6):
+        if x.size[0] <= pixels or x.size[1] <= pixels: return False
+        else: return True
+
+    def resize_multiscales(self, img, borderType=cv2.BORDER_CONSTANT): 
+        def _resize_ratio(img, ratio, fix_h=True):
+            if ratio * self.img_w < self.img_h:
+                if fix_h: trg_h = self.img_h
+                else: trg_h = int(ratio * self.img_w)
+                trg_w = self.img_w
+            else: trg_h, trg_w = self.img_h, int(self.img_h / ratio)
+            img = cv2.resize(img, (trg_w, trg_h))
+            pad_h, pad_w = (self.img_h - trg_h) / 2, (self.img_w - trg_w) / 2
+            top, bottom = math.ceil(pad_h), math.floor(pad_h)
+            left, right = math.ceil(pad_w), math.floor(pad_w)
+            img = cv2.copyMakeBorder(img, top, bottom, left, right, borderType)
+            return img
+        
+        if self.is_training: 
+            if random.random() < 0.5:
+                base, maxh, maxw = self.img_h, self.img_h, self.img_w
+                h, w = random.randint(base, maxh), random.randint(base, maxw)
+                return _resize_ratio(img, h/w)
+            else: return _resize_ratio(img, img.shape[0] / img.shape[1])  # keep aspect ratio
+        else:  return _resize_ratio(img, img.shape[0] / img.shape[1])  # keep aspect ratio
+
+    def resize(self, img):
+        if self.multiscales: return self.resize_multiscales(img, cv2.BORDER_REPLICATE)
+        else: return cv2.resize(img, (self.img_w, self.img_h))
+         
+    def get(self, idx):
+        with self.env.begin(write=False) as txn:
+            image_key, label_key = f'image-{idx+1:09d}', f'label-{idx+1:09d}'
+            try:
+                label = str(txn.get(label_key.encode()), 'utf-8')  # label
+                label = re.sub('[^0-9a-zA-Z]+', '', label)
+                if self.check_length and self.max_length > 0:
+                    if len(label) > self.max_length or len(label) <= 0:
+                        #logging.info(f'Long or short text image is found: {self.name}, {idx}, {label}, {len(label)}')
+                        return self._next_image(idx)
+                label = label[:self.max_length]
+
+                imgbuf = txn.get(image_key.encode())  # image
+                buf = six.BytesIO()
+                buf.write(imgbuf)
+                buf.seek(0)
+                with warnings.catch_warnings():
+                    warnings.simplefilter("ignore", UserWarning) # EXIF warning from TiffPlugin
+                    image = PIL.Image.open(buf).convert(self.convert_mode)
+                if self.is_training and not self._check_image(image):
+                    #logging.info(f'Invalid image is found: {self.name}, {idx}, {label}, {len(label)}')
+                    return self._next_image(idx)
+            except:
+                import traceback
+                traceback.print_exc()
+                logging.info(f'Corrupted image is found: {self.name}, {idx}, {label}, {len(label)}')
+                return self._next_image(idx)
+            return image, label, idx
+
+    def _process_training(self, image):
+        if self.data_aug: image = self.augment_tfs(image)
+        image = self.resize(np.array(image))
+        return image
+
+    def _process_test(self, image):
+        return self.resize(np.array(image)) # TODO:move is_training to here
+
+    def __getitem__(self, idx):
+        image, text, idx_new = self.get(idx)
+        if not self.is_training: assert idx == idx_new, f'idx {idx} != idx_new {idx_new} during testing.'
+
+        if self.is_training: image = self._process_training(image)
+        else: image = self._process_test(image)
+        if self.return_raw: return image, text
+        image = self.totensor(image)
+
+        length = tensor(len(text) + 1).to(dtype=torch.long)  # one for end token
+        label = self.charset.get_labels(text, case_sensitive=self.case_sensitive)
+        label = tensor(label).to(dtype=torch.long)
+        if self.one_hot_y: label = onehot(label, self.charset.num_classes)
+
+        if self.return_idx: y = [label, length, idx_new]
+        else: y = [label, length]
+        return image, y
+
+
+class TextDataset(Dataset):
+    def __init__(self,
+                 path:PathOrStr, 
+                 delimiter:str='\t',
+                 max_length:int=25, 
+                 charset_path:str='data/charset_36.txt', 
+                 case_sensitive=False, 
+                 one_hot_x=True,
+                 one_hot_y=True,
+                 is_training=True,
+                 smooth_label=False,
+                 smooth_factor=0.2,
+                 use_sm=False,
+                 **kwargs):
+        self.path = Path(path)
+        self.case_sensitive, self.use_sm = case_sensitive, use_sm
+        self.smooth_factor, self.smooth_label = smooth_factor, smooth_label
+        self.charset = CharsetMapper(charset_path, max_length=max_length+1)
+        self.one_hot_x, self.one_hot_y, self.is_training = one_hot_x, one_hot_y, is_training
+        if self.is_training and self.use_sm: self.sm = SpellingMutation(charset=self.charset)
+
+        dtype = {'inp': str, 'gt': str}
+        self.df = pd.read_csv(self.path, dtype=dtype, delimiter=delimiter, na_filter=False)
+        self.inp_col, self.gt_col = 0, 1
+
+    def __len__(self): return len(self.df)
+
+    def __getitem__(self, idx):
+        text_x = self.df.iloc[idx, self.inp_col]
+        text_x = re.sub('[^0-9a-zA-Z]+', '', text_x)
+        if not self.case_sensitive: text_x = text_x.lower()
+        if self.is_training and self.use_sm: text_x = self.sm(text_x)
+
+        length_x = tensor(len(text_x) + 1).to(dtype=torch.long)  # one for end token
+        label_x = self.charset.get_labels(text_x, case_sensitive=self.case_sensitive)
+        label_x = tensor(label_x)
+        if self.one_hot_x:
+            label_x = onehot(label_x, self.charset.num_classes)
+            if self.is_training and self.smooth_label: 
+                label_x = torch.stack([self.prob_smooth_label(l) for l in label_x])
+        x =  [label_x, length_x]
+    
+        text_y = self.df.iloc[idx, self.gt_col]
+        text_y = re.sub('[^0-9a-zA-Z]+', '', text_y)
+        if not self.case_sensitive: text_y = text_y.lower()
+        length_y = tensor(len(text_y) + 1).to(dtype=torch.long)  # one for end token
+        label_y = self.charset.get_labels(text_y, case_sensitive=self.case_sensitive)
+        label_y = tensor(label_y)
+        if self.one_hot_y: label_y = onehot(label_y, self.charset.num_classes)
+        y = [label_y, length_y]
+
+        return x, y
+
+    def prob_smooth_label(self, one_hot):
+        one_hot = one_hot.float()
+        delta = torch.rand([]) * self.smooth_factor
+        num_classes = len(one_hot)
+        noise = torch.rand(num_classes)
+        noise = noise / noise.sum() * delta
+        one_hot = one_hot * (1 - delta) + noise
+        return one_hot
+
+
+class SpellingMutation(object):
+    def __init__(self, pn0=0.7, pn1=0.85, pn2=0.95, pt0=0.7, pt1=0.85, charset=None):
+        """ 
+        Args:
+            pn0: the prob of not modifying characters is (pn0)
+            pn1: the prob of modifying one characters is (pn1 - pn0)
+            pn2: the prob of modifying two characters is (pn2 - pn1), 
+                 and three (1 - pn2)
+            pt0: the prob of replacing operation is pt0.
+            pt1: the prob of inserting operation is (pt1 - pt0),
+                 and deleting operation is (1 - pt1)
+        """
+        super().__init__()
+        self.pn0, self.pn1, self.pn2 = pn0, pn1, pn2
+        self.pt0, self.pt1 = pt0, pt1
+        self.charset = charset
+        logging.info(f'the probs: pn0={self.pn0}, pn1={self.pn1} ' + 
+                     f'pn2={self.pn2}, pt0={self.pt0}, pt1={self.pt1}')
+
+    def is_digit(self, text, ratio=0.5):
+        length = max(len(text), 1)
+        digit_num = sum([t in self.charset.digits for t in text])
+        if digit_num / length < ratio: return False
+        return True
+
+    def is_unk_char(self, char):
+        # return char == self.charset.unk_char
+        return (char not in self.charset.digits) and (char not in self.charset.alphabets)
+
+    def get_num_to_modify(self, length):
+        prob = random.random()
+        if prob < self.pn0: num_to_modify = 0
+        elif prob < self.pn1: num_to_modify = 1
+        elif prob < self.pn2: num_to_modify = 2
+        else: num_to_modify = 3
+        
+        if length <= 1: num_to_modify = 0
+        elif length >= 2 and length <= 4: num_to_modify = min(num_to_modify, 1)
+        else: num_to_modify = min(num_to_modify, length // 2)  # smaller than length // 2
+        return num_to_modify
+
+    def __call__(self, text, debug=False):
+        if self.is_digit(text): return text
+        length = len(text)
+        num_to_modify = self.get_num_to_modify(length)
+        if num_to_modify <= 0: return text
+
+        chars = []
+        index = np.arange(0, length)
+        random.shuffle(index)
+        index = index[: num_to_modify]
+        if debug: self.index = index
+        for i, t in enumerate(text):
+            if i not in index: chars.append(t)
+            elif self.is_unk_char(t): chars.append(t)
+            else:
+                prob = random.random()
+                if prob < self.pt0: # replace
+                    chars.append(random.choice(self.charset.alphabets))
+                elif prob < self.pt1: # insert
+                    chars.append(random.choice(self.charset.alphabets))
+                    chars.append(t)
+                else: # delete
+                    continue
+        new_text = ''.join(chars[: self.charset.max_length-1])
+        return new_text if len(new_text) >= 1 else text

demo.py

@@ -0,0 +1,109 @@
+import argparse
+import logging
+import os
+import glob
+import tqdm
+import torch
+import PIL
+import cv2
+import numpy as np
+import torch.nn.functional as F
+from torchvision import transforms
+from utils import Config, Logger, CharsetMapper
+
+def get_model(config):
+    import importlib
+    names = config.model_name.split('.')
+    module_name, class_name = '.'.join(names[:-1]), names[-1]
+    cls = getattr(importlib.import_module(module_name), class_name)
+    model = cls(config)
+    logging.info(model)
+    model = model.eval()
+    return model
+
+def preprocess(img, width, height):
+    img = cv2.resize(np.array(img), (width, height))
+    img = transforms.ToTensor()(img).unsqueeze(0)
+    mean = torch.tensor([0.485, 0.456, 0.406])
+    std  = torch.tensor([0.229, 0.224, 0.225])
+    return (img-mean[...,None,None]) / std[...,None,None]
+
+def postprocess(output, charset, model_eval):
+    def _get_output(last_output, model_eval):
+        if isinstance(last_output, (tuple, list)): 
+            for res in last_output:
+                if res['name'] == model_eval: output = res
+        else: output = last_output
+        return output
+
+    def _decode(logit):
+        """ Greed decode """
+        out = F.softmax(logit, dim=2)
+        pt_text, pt_scores, pt_lengths = [], [], []
+        for o in out:
+            text = charset.get_text(o.argmax(dim=1), padding=False, trim=False)
+            text = text.split(charset.null_char)[0]  # end at end-token
+            pt_text.append(text)
+            pt_scores.append(o.max(dim=1)[0])
+            pt_lengths.append(min(len(text) + 1, charset.max_length))  # one for end-token
+        return pt_text, pt_scores, pt_lengths
+
+    output = _get_output(output, model_eval)
+    logits, pt_lengths = output['logits'], output['pt_lengths']
+    pt_text, pt_scores, pt_lengths_ = _decode(logits)
+    
+    return pt_text, pt_scores, pt_lengths_
+
+def load(model, file, device=None, strict=True):
+    if device is None: device = 'cpu'
+    elif isinstance(device, int): device = torch.device('cuda', device)
+    assert os.path.isfile(file)
+    state = torch.load(file, map_location=device)
+    if set(state.keys()) == {'model', 'opt'}:
+        state = state['model']
+    model.load_state_dict(state, strict=strict)
+    return model
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--config', type=str, default='configs/train_abinet.yaml',
+                        help='path to config file')
+    parser.add_argument('--input', type=str, default='figs/test')
+    parser.add_argument('--cuda', type=int, default=-1)
+    parser.add_argument('--checkpoint', type=str, default='workdir/train-abinet/best-train-abinet.pth')
+    parser.add_argument('--model_eval', type=str, default='alignment', 
+                        choices=['alignment', 'vision', 'language'])
+    args = parser.parse_args()
+    config = Config(args.config)
+    if args.checkpoint is not None: config.model_checkpoint = args.checkpoint
+    if args.model_eval is not None: config.model_eval = args.model_eval
+    config.global_phase = 'test'
+    config.model_vision_checkpoint, config.model_language_checkpoint = None, None
+    device = 'cpu' if args.cuda < 0 else f'cuda:{args.cuda}'
+
+    Logger.init(config.global_workdir, config.global_name, config.global_phase)
+    Logger.enable_file()
+    logging.info(config)
+
+    logging.info('Construct model.')
+    model = get_model(config).to(device)
+    model = load(model, config.model_checkpoint, device=device)
+    charset = CharsetMapper(filename=config.dataset_charset_path,
+                            max_length=config.dataset_max_length + 1)
+
+    if os.path.isdir(args.input):
+        paths = [os.path.join(args.input, fname) for fname in os.listdir(args.input)]
+    else:
+        paths = glob.glob(os.path.expanduser(args.input))
+        assert paths, "The input path(s) was not found"
+    paths = sorted(paths)
+    for path in tqdm.tqdm(paths):
+        img = PIL.Image.open(path).convert('RGB')
+        img = preprocess(img, config.dataset_image_width, config.dataset_image_height)
+        img = img.to(device)
+        res = model(img)
+        pt_text, _, __ = postprocess(res, charset, config.model_eval)
+        logging.info(f'{path}: {pt_text[0]}')
+
+if __name__ == '__main__':
+    main()

docker/Dockerfile

@@ -0,0 +1,25 @@
+FROM anibali/pytorch:cuda-9.0
+MAINTAINER fangshancheng <[email protected]>
+RUN sudo rm -rf /etc/apt/sources.list.d && \
+    sudo apt update && \
+    sudo apt install -y build-essential vim && \
+    conda config --add channels https://mirrors.ustc.edu.cn/anaconda/pkgs/free/ && \
+    conda config --add channels https://mirrors.ustc.edu.cn/anaconda/pkgs/main/ && \
+    conda config --set show_channel_urls yes && \
+    pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/ && \
+    pip install torch==1.1.0 torchvision==0.3.0 && \
+    pip install fastai==1.0.60 && \
+    pip install ipdb jupyter ipython lmdb editdistance tensorboardX natsort nltk && \
+    conda uninstall -y --force pillow pil jpeg libtiff libjpeg-turbo && \
+    pip uninstall -y pillow pil jpeg libtiff libjpeg-turbo && \
+    conda install -yc conda-forge libjpeg-turbo && \
+    CFLAGS="${CFLAGS} -mavx2" pip install --no-cache-dir --force-reinstall --no-binary :all: --compile pillow-simd==6.2.2.post1 && \
+    conda install -y jpeg libtiff opencv && \
+    sudo rm -rf /var/lib/apt/lists/* && \
+    sudo rm -rf /tmp/* && \
+    sudo rm -rf ~/.cache && \
+    sudo apt clean all && \
+    conda clean -y -a
+EXPOSE 8888
+ENV LANG C.UTF-8
+ENV LC_ALL C.UTF-8

losses.py

@@ -0,0 +1,72 @@
+from fastai.vision import *
+
+from modules.model import Model
+
+
+class MultiLosses(nn.Module):
+    def __init__(self, one_hot=True):
+        super().__init__()
+        self.ce = SoftCrossEntropyLoss() if one_hot else torch.nn.CrossEntropyLoss()
+        self.bce = torch.nn.BCELoss()
+
+    @property
+    def last_losses(self):
+        return self.losses
+
+    def _flatten(self, sources, lengths):
+        return torch.cat([t[:l] for t, l in zip(sources, lengths)])
+
+    def _merge_list(self, all_res):
+        if not isinstance(all_res, (list, tuple)):
+            return all_res
+        def merge(items):
+            if isinstance(items[0], torch.Tensor): return torch.cat(items, dim=0)
+            else: return items[0]
+        res = dict()
+        for key in all_res[0].keys():
+            items = [r[key] for r in all_res]
+            res[key] = merge(items)
+        return res
+
+    def _ce_loss(self, output, gt_labels, gt_lengths, idx=None, record=True):
+        loss_name = output.get('name')
+        pt_logits, weight = output['logits'], output['loss_weight']
+
+        assert pt_logits.shape[0] % gt_labels.shape[0] == 0
+        iter_size = pt_logits.shape[0] // gt_labels.shape[0]
+        if iter_size > 1:
+            gt_labels = gt_labels.repeat(3, 1, 1)
+            gt_lengths = gt_lengths.repeat(3)
+        flat_gt_labels = self._flatten(gt_labels, gt_lengths)
+        flat_pt_logits = self._flatten(pt_logits, gt_lengths)
+
+        nll = output.get('nll')
+        if nll is not None:
+            loss = self.ce(flat_pt_logits, flat_gt_labels, softmax=False) * weight
+        else:
+            loss = self.ce(flat_pt_logits, flat_gt_labels) * weight
+        if record and loss_name is not None: self.losses[f'{loss_name}_loss'] = loss
+
+        return loss
+
+    def forward(self, outputs, *args):
+        self.losses = {}
+        if isinstance(outputs, (tuple, list)):
+            outputs = [self._merge_list(o) for o in outputs]
+            return sum([self._ce_loss(o, *args) for o in outputs if o['loss_weight'] > 0.])
+        else:
+            return self._ce_loss(outputs, *args, record=False)
+
+
+class SoftCrossEntropyLoss(nn.Module):
+    def __init__(self, reduction="mean"):
+        super().__init__()
+        self.reduction = reduction
+
+    def forward(self, input, target, softmax=True):
+        if softmax: log_prob = F.log_softmax(input, dim=-1)
+        else: log_prob = torch.log(input)
+        loss = -(target * log_prob).sum(dim=-1)
+        if self.reduction == "mean": return loss.mean()
+        elif self.reduction == "sum": return loss.sum()
+        else: return loss

main.py

@@ -0,0 +1,246 @@
+import argparse
+import logging
+import os
+import random
+
+import torch
+from fastai.callbacks.general_sched import GeneralScheduler, TrainingPhase
+from fastai.distributed import *
+from fastai.vision import *
+from torch.backends import cudnn
+
+from callbacks import DumpPrediction, IterationCallback, TextAccuracy, TopKTextAccuracy
+from dataset import ImageDataset, TextDataset
+from losses import MultiLosses
+from utils import Config, Logger, MyDataParallel, MyConcatDataset
+
+
+def _set_random_seed(seed):
+    if seed is not None:
+        random.seed(seed)
+        torch.manual_seed(seed)
+        cudnn.deterministic = True
+        logging.warning('You have chosen to seed training. '
+                        'This will slow down your training!')
+
+def _get_training_phases(config, n):
+    lr = np.array(config.optimizer_lr)
+    periods = config.optimizer_scheduler_periods
+    sigma = [config.optimizer_scheduler_gamma ** i for i in range(len(periods))]
+    phases = [TrainingPhase(n * periods[i]).schedule_hp('lr', lr * sigma[i])
+                for i in range(len(periods))]
+    return phases
+
+def _get_dataset(ds_type, paths, is_training, config, **kwargs):
+    kwargs.update({
+        'img_h': config.dataset_image_height,
+        'img_w': config.dataset_image_width,
+        'max_length': config.dataset_max_length,
+        'case_sensitive': config.dataset_case_sensitive,
+        'charset_path': config.dataset_charset_path,
+        'data_aug': config.dataset_data_aug,
+        'deteriorate_ratio': config.dataset_deteriorate_ratio,
+        'is_training': is_training,
+        'multiscales': config.dataset_multiscales,
+        'one_hot_y': config.dataset_one_hot_y,
+    })
+    datasets = [ds_type(p, **kwargs) for p in paths]
+    if len(datasets) > 1: return MyConcatDataset(datasets)
+    else: return datasets[0]
+
+
+def _get_language_databaunch(config):
+    kwargs = {
+        'max_length': config.dataset_max_length,
+        'case_sensitive': config.dataset_case_sensitive,
+        'charset_path': config.dataset_charset_path,
+        'smooth_label': config.dataset_smooth_label,
+        'smooth_factor': config.dataset_smooth_factor,
+        'one_hot_y': config.dataset_one_hot_y,
+        'use_sm': config.dataset_use_sm,
+    }
+    train_ds = TextDataset(config.dataset_train_roots[0], is_training=True, **kwargs)
+    valid_ds = TextDataset(config.dataset_test_roots[0], is_training=False, **kwargs)
+    data = DataBunch.create(
+        path=train_ds.path,
+        train_ds=train_ds,
+        valid_ds=valid_ds,
+        bs=config.dataset_train_batch_size,
+        val_bs=config.dataset_test_batch_size,
+        num_workers=config.dataset_num_workers,
+        pin_memory=config.dataset_pin_memory)
+    logging.info(f'{len(data.train_ds)} training items found.')
+    if not data.empty_val:
+        logging.info(f'{len(data.valid_ds)} valid items found.')
+    return data
+
+def _get_databaunch(config):
+    # An awkward way to reduce loadding data time during test
+    if config.global_phase == 'test': config.dataset_train_roots = config.dataset_test_roots
+    train_ds = _get_dataset(ImageDataset, config.dataset_train_roots, True, config)
+    valid_ds = _get_dataset(ImageDataset, config.dataset_test_roots, False, config)
+    data = ImageDataBunch.create(
+        train_ds=train_ds,
+        valid_ds=valid_ds,
+        bs=config.dataset_train_batch_size,
+        val_bs=config.dataset_test_batch_size,
+        num_workers=config.dataset_num_workers,
+        pin_memory=config.dataset_pin_memory).normalize(imagenet_stats)
+    ar_tfm = lambda x: ((x[0], x[1]), x[1])  # auto-regression only for dtd
+    data.add_tfm(ar_tfm)
+
+    logging.info(f'{len(data.train_ds)} training items found.')
+    if not data.empty_val:
+        logging.info(f'{len(data.valid_ds)} valid items found.')
+    
+    return data
+
+def _get_model(config):
+    import importlib
+    names = config.model_name.split('.')
+    module_name, class_name = '.'.join(names[:-1]), names[-1]
+    cls = getattr(importlib.import_module(module_name), class_name)
+    model = cls(config)
+    logging.info(model)
+    return model
+
+
+def _get_learner(config, data, model, local_rank=None):
+    strict = ifnone(config.model_strict, True)
+    if config.global_stage == 'pretrain-language':
+        metrics = [TopKTextAccuracy(
+            k=ifnone(config.model_k, 5),
+            charset_path=config.dataset_charset_path,
+            max_length=config.dataset_max_length + 1,
+            case_sensitive=config.dataset_eval_case_sensisitves,
+            model_eval=config.model_eval)] 
+    else:
+        metrics = [TextAccuracy(
+            charset_path=config.dataset_charset_path,
+            max_length=config.dataset_max_length + 1,
+            case_sensitive=config.dataset_eval_case_sensisitves,
+            model_eval=config.model_eval)]
+    opt_type = getattr(torch.optim, config.optimizer_type)
+    learner = Learner(data, model, silent=True, model_dir='.',
+        true_wd=config.optimizer_true_wd, 
+        wd=config.optimizer_wd,
+        bn_wd=config.optimizer_bn_wd,
+        path=config.global_workdir,
+        metrics=metrics,
+        opt_func=partial(opt_type, **config.optimizer_args or dict()), 
+        loss_func=MultiLosses(one_hot=config.dataset_one_hot_y))
+    learner.split(lambda m: children(m))
+
+    if config.global_phase == 'train':
+        num_replicas = 1 if local_rank is None else torch.distributed.get_world_size()
+        phases = _get_training_phases(config, len(learner.data.train_dl)//num_replicas)
+        learner.callback_fns += [
+            partial(GeneralScheduler, phases=phases),
+            partial(GradientClipping, clip=config.optimizer_clip_grad),
+            partial(IterationCallback, name=config.global_name,
+                    show_iters=config.training_show_iters,
+                    eval_iters=config.training_eval_iters,
+                    save_iters=config.training_save_iters,
+                    start_iters=config.training_start_iters,
+                    stats_iters=config.training_stats_iters)]
+    else:
+        learner.callbacks += [
+            DumpPrediction(learn=learner,
+                    dataset='-'.join([Path(p).name for p in config.dataset_test_roots]),charset_path=config.dataset_charset_path,
+                    model_eval=config.model_eval,
+                    debug=config.global_debug,
+                    image_only=config.global_image_only)]
+
+    learner.rank = local_rank
+    if local_rank is not None:
+        logging.info(f'Set model to distributed with rank {local_rank}.')
+        learner.model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(learner.model)
+        learner.model.to(local_rank)
+        learner = learner.to_distributed(local_rank)
+
+    if torch.cuda.device_count() > 1 and local_rank is None:
+        logging.info(f'Use {torch.cuda.device_count()} GPUs.')
+        learner.model = MyDataParallel(learner.model)
+
+    if config.model_checkpoint:
+        if Path(config.model_checkpoint).exists():
+            with open(config.model_checkpoint, 'rb') as f:
+                buffer = io.BytesIO(f.read())
+            learner.load(buffer, strict=strict)
+        else:
+            from distutils.dir_util import copy_tree
+            src = Path('/data/fangsc/model')/config.global_name
+            trg = Path('/output')/config.global_name
+            if src.exists(): copy_tree(str(src), str(trg))
+            learner.load(config.model_checkpoint, strict=strict)
+        logging.info(f'Read model from {config.model_checkpoint}')
+    elif config.global_phase == 'test':
+        learner.load(f'best-{config.global_name}', strict=strict)
+        logging.info(f'Read model from best-{config.global_name}')
+
+    if learner.opt_func.func.__name__ == 'Adadelta':    # fastai bug, fix after 1.0.60
+        learner.fit(epochs=0, lr=config.optimizer_lr)
+        learner.opt.mom = 0.
+
+    return learner
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--config', type=str, required=True,
+                        help='path to config file')
+    parser.add_argument('--phase', type=str, default=None, choices=['train', 'test'])
+    parser.add_argument('--name', type=str, default=None)
+    parser.add_argument('--checkpoint', type=str, default=None)
+    parser.add_argument('--test_root', type=str, default=None)
+    parser.add_argument("--local_rank", type=int, default=None)
+    parser.add_argument('--debug', action='store_true', default=None)
+    parser.add_argument('--image_only', action='store_true', default=None)
+    parser.add_argument('--model_strict', action='store_false', default=None)
+    parser.add_argument('--model_eval', type=str, default=None, 
+                        choices=['alignment', 'vision', 'language'])
+    args = parser.parse_args()
+    config = Config(args.config)
+    if args.name is not None: config.global_name = args.name
+    if args.phase is not None: config.global_phase = args.phase
+    if args.test_root is not None: config.dataset_test_roots = [args.test_root]
+    if args.checkpoint is not None: config.model_checkpoint = args.checkpoint
+    if args.debug is not None: config.global_debug = args.debug
+    if args.image_only is not None: config.global_image_only = args.image_only
+    if args.model_eval is not None: config.model_eval = args.model_eval
+    if args.model_strict is not None: config.model_strict = args.model_strict
+
+    Logger.init(config.global_workdir, config.global_name, config.global_phase)
+    Logger.enable_file()
+    _set_random_seed(config.global_seed)
+    logging.info(config)
+
+    if args.local_rank is not None:
+        logging.info(f'Init distribution training at device {args.local_rank}.')
+        torch.cuda.set_device(args.local_rank)
+        torch.distributed.init_process_group(backend='nccl', init_method='env://')
+
+    logging.info('Construct dataset.')
+    if config.global_stage == 'pretrain-language': data = _get_language_databaunch(config)
+    else: data = _get_databaunch(config)
+
+    logging.info('Construct model.')
+    model = _get_model(config)
+
+    logging.info('Construct learner.')
+    learner = _get_learner(config, data, model, args.local_rank)
+
+    if config.global_phase == 'train':
+        logging.info('Start training.')
+        learner.fit(epochs=config.training_epochs,
+                    lr=config.optimizer_lr)
+    else:
+        logging.info('Start validate')
+        last_metrics = learner.validate()
+        log_str = f'eval loss = {last_metrics[0]:6.3f},  ' \
+                  f'ccr = {last_metrics[1]:6.3f},  cwr = {last_metrics[2]:6.3f},  ' \
+                  f'ted = {last_metrics[3]:6.3f},  ned = {last_metrics[4]:6.0f},  ' \
+                  f'ted/w = {last_metrics[5]:6.3f}, '
+        logging.info(log_str)
+
+if __name__ == '__main__':
+    main()

modules/attention.py

@@ -0,0 +1,97 @@
+import torch
+import torch.nn as nn
+from .transformer import PositionalEncoding
+
+class Attention(nn.Module):
+    def __init__(self, in_channels=512, max_length=25, n_feature=256):
+        super().__init__()
+        self.max_length = max_length
+
+        self.f0_embedding = nn.Embedding(max_length, in_channels)
+        self.w0 = nn.Linear(max_length, n_feature)
+        self.wv = nn.Linear(in_channels, in_channels)
+        self.we = nn.Linear(in_channels, max_length)
+
+        self.active = nn.Tanh()
+        self.softmax = nn.Softmax(dim=2)
+
+    def forward(self, enc_output):
+        enc_output = enc_output.permute(0, 2, 3, 1).flatten(1, 2)
+        reading_order = torch.arange(self.max_length, dtype=torch.long, device=enc_output.device)
+        reading_order = reading_order.unsqueeze(0).expand(enc_output.size(0), -1)  # (S,) -> (B, S)
+        reading_order_embed = self.f0_embedding(reading_order)  # b,25,512
+
+        t = self.w0(reading_order_embed.permute(0, 2, 1))  # b,512,256
+        t = self.active(t.permute(0, 2, 1) + self.wv(enc_output))  # b,256,512
+
+        attn = self.we(t)  # b,256,25
+        attn = self.softmax(attn.permute(0, 2, 1))  # b,25,256
+        g_output = torch.bmm(attn, enc_output)  # b,25,512
+        return g_output, attn.view(*attn.shape[:2], 8, 32)
+
+
+def encoder_layer(in_c, out_c, k=3, s=2, p=1):
+    return nn.Sequential(nn.Conv2d(in_c, out_c, k, s, p),
+                         nn.BatchNorm2d(out_c),
+                         nn.ReLU(True))
+
+def decoder_layer(in_c, out_c, k=3, s=1, p=1, mode='nearest', scale_factor=None, size=None):
+    align_corners = None if mode=='nearest' else True
+    return nn.Sequential(nn.Upsample(size=size, scale_factor=scale_factor, 
+                                     mode=mode, align_corners=align_corners),
+                         nn.Conv2d(in_c, out_c, k, s, p),
+                         nn.BatchNorm2d(out_c),
+                         nn.ReLU(True))
+
+
+class PositionAttention(nn.Module):
+    def __init__(self, max_length, in_channels=512, num_channels=64, 
+                 h=8, w=32, mode='nearest', **kwargs):
+        super().__init__()
+        self.max_length = max_length
+        self.k_encoder = nn.Sequential(
+            encoder_layer(in_channels, num_channels, s=(1, 2)),
+            encoder_layer(num_channels, num_channels, s=(2, 2)),
+            encoder_layer(num_channels, num_channels, s=(2, 2)),
+            encoder_layer(num_channels, num_channels, s=(2, 2))
+        )
+        self.k_decoder = nn.Sequential(
+            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
+            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
+            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
+            decoder_layer(num_channels, in_channels, size=(h, w), mode=mode)
+        )
+
+        self.pos_encoder = PositionalEncoding(in_channels, dropout=0, max_len=max_length)
+        self.project = nn.Linear(in_channels, in_channels)
+
+    def forward(self, x):
+        N, E, H, W = x.size()
+        k, v = x, x  # (N, E, H, W)
+
+        # calculate key vector
+        features = []
+        for i in range(0, len(self.k_encoder)):
+            k = self.k_encoder[i](k)
+            features.append(k)
+        for i in range(0, len(self.k_decoder) - 1):
+            k = self.k_decoder[i](k)
+            k = k + features[len(self.k_decoder) - 2 - i]
+        k = self.k_decoder[-1](k)
+
+        # calculate query vector
+        # TODO q=f(q,k)
+        zeros = x.new_zeros((self.max_length, N, E))  # (T, N, E)
+        q = self.pos_encoder(zeros)  # (T, N, E)
+        q = q.permute(1, 0, 2)  # (N, T, E)
+        q = self.project(q)  # (N, T, E)
+        
+        # calculate attention
+        attn_scores = torch.bmm(q, k.flatten(2, 3))  # (N, T, (H*W))
+        attn_scores = attn_scores / (E ** 0.5)
+        attn_scores = torch.softmax(attn_scores, dim=-1)
+
+        v = v.permute(0, 2, 3, 1).view(N, -1, E)  # (N, (H*W), E)
+        attn_vecs = torch.bmm(attn_scores, v)  # (N, T, E)
+
+        return attn_vecs, attn_scores.view(N, -1, H, W)

modules/backbone.py

@@ -0,0 +1,36 @@
+import torch
+import torch.nn as nn
+from fastai.vision import *
+
+from modules.model import _default_tfmer_cfg
+from modules.resnet import resnet45
+from modules.transformer import (PositionalEncoding,
+                                 TransformerEncoder,
+                                 TransformerEncoderLayer)
+
+
+class ResTranformer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.resnet = resnet45()
+
+        self.d_model = ifnone(config.model_vision_d_model, _default_tfmer_cfg['d_model'])
+        nhead = ifnone(config.model_vision_nhead, _default_tfmer_cfg['nhead'])
+        d_inner = ifnone(config.model_vision_d_inner, _default_tfmer_cfg['d_inner'])
+        dropout = ifnone(config.model_vision_dropout, _default_tfmer_cfg['dropout'])
+        activation = ifnone(config.model_vision_activation, _default_tfmer_cfg['activation'])
+        num_layers = ifnone(config.model_vision_backbone_ln, 2)
+
+        self.pos_encoder = PositionalEncoding(self.d_model, max_len=8*32)
+        encoder_layer = TransformerEncoderLayer(d_model=self.d_model, nhead=nhead, 
+                dim_feedforward=d_inner, dropout=dropout, activation=activation)
+        self.transformer = TransformerEncoder(encoder_layer, num_layers)
+
+    def forward(self, images):
+        feature = self.resnet(images)
+        n, c, h, w = feature.shape
+        feature = feature.view(n, c, -1).permute(2, 0, 1)
+        feature = self.pos_encoder(feature)
+        feature = self.transformer(feature)
+        feature = feature.permute(1, 2, 0).view(n, c, h, w)
+        return feature

modules/model.py

@@ -0,0 +1,50 @@
+import torch
+import torch.nn as nn
+
+from utils import CharsetMapper
+
+
+_default_tfmer_cfg = dict(d_model=512, nhead=8, d_inner=2048, # 1024
+                          dropout=0.1, activation='relu')
+
+class Model(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.max_length = config.dataset_max_length + 1
+        self.charset = CharsetMapper(config.dataset_charset_path, max_length=self.max_length)
+    
+    def load(self, source, device=None, strict=True):
+        state = torch.load(source, map_location=device)
+        self.load_state_dict(state['model'], strict=strict)
+
+    def _get_length(self, logit, dim=-1):
+        """ Greed decoder to obtain length from logit"""
+        out = (logit.argmax(dim=-1) == self.charset.null_label)
+        abn = out.any(dim)
+        out = ((out.cumsum(dim) == 1) & out).max(dim)[1]
+        out = out + 1  # additional end token
+        out = torch.where(abn, out, out.new_tensor(logit.shape[1]))
+        return out
+
+    @staticmethod
+    def _get_padding_mask(length, max_length):
+        length = length.unsqueeze(-1)
+        grid = torch.arange(0, max_length, device=length.device).unsqueeze(0)
+        return grid >= length
+
+    @staticmethod
+    def _get_square_subsequent_mask(sz, device, diagonal=0, fw=True):
+        r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf').
+            Unmasked positions are filled with float(0.0).
+        """
+        mask = (torch.triu(torch.ones(sz, sz, device=device), diagonal=diagonal) == 1)
+        if fw: mask = mask.transpose(0, 1)
+        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+        return mask
+
+    @staticmethod
+    def _get_location_mask(sz, device=None):
+        mask = torch.eye(sz, device=device)
+        mask = mask.float().masked_fill(mask == 1, float('-inf'))
+        return mask

modules/model_abinet.py

@@ -0,0 +1,30 @@
+import torch
+import torch.nn as nn
+from fastai.vision import *
+
+from .model_vision import BaseVision
+from .model_language import BCNLanguage
+from .model_alignment import BaseAlignment
+
+
+class ABINetModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.use_alignment = ifnone(config.model_use_alignment, True)
+        self.max_length = config.dataset_max_length + 1  # additional stop token
+        self.vision = BaseVision(config)
+        self.language = BCNLanguage(config)
+        if self.use_alignment: self.alignment = BaseAlignment(config)
+
+    def forward(self, images, *args):
+        v_res = self.vision(images)
+        v_tokens = torch.softmax(v_res['logits'], dim=-1)
+        v_lengths = v_res['pt_lengths'].clamp_(2, self.max_length)  # TODO:move to langauge model
+
+        l_res = self.language(v_tokens, v_lengths)
+        if not self.use_alignment:
+            return l_res, v_res
+        l_feature, v_feature = l_res['feature'], v_res['feature']
+        
+        a_res = self.alignment(l_feature, v_feature)
+        return a_res, l_res, v_res

modules/model_abinet_iter.py

@@ -0,0 +1,34 @@
+import torch
+import torch.nn as nn
+from fastai.vision import *
+
+from .model_vision import BaseVision
+from .model_language import BCNLanguage
+from .model_alignment import BaseAlignment
+
+
+class ABINetIterModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.iter_size = ifnone(config.model_iter_size, 1)
+        self.max_length = config.dataset_max_length + 1  # additional stop token
+        self.vision = BaseVision(config)
+        self.language = BCNLanguage(config)
+        self.alignment = BaseAlignment(config)
+
+    def forward(self, images, *args):
+        v_res = self.vision(images)
+        a_res = v_res
+        all_l_res, all_a_res = [], []
+        for _ in range(self.iter_size):
+            tokens = torch.softmax(a_res['logits'], dim=-1)
+            lengths = a_res['pt_lengths']
+            lengths.clamp_(2, self.max_length)  # TODO:move to langauge model
+            l_res = self.language(tokens, lengths)
+            all_l_res.append(l_res)
+            a_res = self.alignment(l_res['feature'], v_res['feature'])
+            all_a_res.append(a_res)
+        if self.training:
+            return all_a_res, all_l_res, v_res
+        else:
+            return a_res, all_l_res[-1], v_res

modules/model_alignment.py

@@ -0,0 +1,34 @@
+import torch
+import torch.nn as nn
+from fastai.vision import *
+
+from modules.model import Model, _default_tfmer_cfg
+
+
+class BaseAlignment(Model):
+    def __init__(self, config):
+        super().__init__(config)
+        d_model = ifnone(config.model_alignment_d_model, _default_tfmer_cfg['d_model'])
+
+        self.loss_weight = ifnone(config.model_alignment_loss_weight, 1.0)
+        self.max_length = config.dataset_max_length + 1  # additional stop token
+        self.w_att = nn.Linear(2 * d_model, d_model)
+        self.cls = nn.Linear(d_model, self.charset.num_classes)
+
+    def forward(self, l_feature, v_feature):
+        """
+        Args:
+            l_feature: (N, T, E) where T is length, N is batch size and d is dim of model
+            v_feature: (N, T, E) shape the same as l_feature 
+            l_lengths: (N,)
+            v_lengths: (N,)
+        """
+        f = torch.cat((l_feature, v_feature), dim=2)
+        f_att = torch.sigmoid(self.w_att(f))
+        output = f_att * v_feature + (1 - f_att) * l_feature
+
+        logits = self.cls(output)  # (N, T, C)
+        pt_lengths = self._get_length(logits)
+
+        return {'logits': logits, 'pt_lengths': pt_lengths, 'loss_weight':self.loss_weight,
+                'name': 'alignment'}

modules/model_language.py

@@ -0,0 +1,67 @@
+import logging
+import torch.nn as nn
+from fastai.vision import *
+
+from modules.model import _default_tfmer_cfg
+from modules.model import Model
+from modules.transformer import (PositionalEncoding, 
+                                 TransformerDecoder,
+                                 TransformerDecoderLayer)
+
+
+class BCNLanguage(Model):
+    def __init__(self, config):
+        super().__init__(config)
+        d_model = ifnone(config.model_language_d_model, _default_tfmer_cfg['d_model'])
+        nhead = ifnone(config.model_language_nhead, _default_tfmer_cfg['nhead'])
+        d_inner = ifnone(config.model_language_d_inner, _default_tfmer_cfg['d_inner'])
+        dropout = ifnone(config.model_language_dropout, _default_tfmer_cfg['dropout'])
+        activation = ifnone(config.model_language_activation, _default_tfmer_cfg['activation'])
+        num_layers = ifnone(config.model_language_num_layers, 4)
+        self.d_model = d_model
+        self.detach = ifnone(config.model_language_detach, True)
+        self.use_self_attn = ifnone(config.model_language_use_self_attn, False)
+        self.loss_weight = ifnone(config.model_language_loss_weight, 1.0)
+        self.max_length = config.dataset_max_length + 1  # additional stop token
+        self.debug = ifnone(config.global_debug, False)
+
+        self.proj = nn.Linear(self.charset.num_classes, d_model, False)
+        self.token_encoder = PositionalEncoding(d_model, max_len=self.max_length)
+        self.pos_encoder = PositionalEncoding(d_model, dropout=0, max_len=self.max_length)
+        decoder_layer = TransformerDecoderLayer(d_model, nhead, d_inner, dropout, 
+                activation, self_attn=self.use_self_attn, debug=self.debug)
+        self.model = TransformerDecoder(decoder_layer, num_layers)
+
+        self.cls = nn.Linear(d_model, self.charset.num_classes)
+
+        if config.model_language_checkpoint is not None:
+            logging.info(f'Read language model from {config.model_language_checkpoint}.')
+            self.load(config.model_language_checkpoint)
+
+    def forward(self, tokens, lengths):
+        """
+        Args:
+            tokens: (N, T, C) where T is length, N is batch size and C is classes number
+            lengths: (N,)
+        """
+        if self.detach: tokens = tokens.detach()
+        embed = self.proj(tokens)  # (N, T, E)
+        embed = embed.permute(1, 0, 2)  # (T, N, E)
+        embed = self.token_encoder(embed)  # (T, N, E)
+        padding_mask = self._get_padding_mask(lengths, self.max_length)
+
+        zeros = embed.new_zeros(*embed.shape)
+        qeury = self.pos_encoder(zeros)
+        location_mask = self._get_location_mask(self.max_length, tokens.device)
+        output = self.model(qeury, embed,
+                tgt_key_padding_mask=padding_mask,
+                memory_mask=location_mask,
+                memory_key_padding_mask=padding_mask)  # (T, N, E)
+        output = output.permute(1, 0, 2)  # (N, T, E)
+
+        logits = self.cls(output)  # (N, T, C)
+        pt_lengths = self._get_length(logits)
+
+        res =  {'feature': output, 'logits': logits, 'pt_lengths': pt_lengths,
+                'loss_weight':self.loss_weight, 'name': 'language'}
+        return res

modules/model_vision.py

@@ -0,0 +1,47 @@
+import logging
+import torch.nn as nn
+from fastai.vision import *
+
+from modules.attention import *
+from modules.backbone import ResTranformer
+from modules.model import Model
+from modules.resnet import resnet45
+
+
+class BaseVision(Model):
+    def __init__(self, config):
+        super().__init__(config)
+        self.loss_weight = ifnone(config.model_vision_loss_weight, 1.0)
+        self.out_channels = ifnone(config.model_vision_d_model, 512)
+
+        if config.model_vision_backbone == 'transformer':
+            self.backbone = ResTranformer(config)
+        else: self.backbone = resnet45()
+        
+        if config.model_vision_attention == 'position':
+            mode = ifnone(config.model_vision_attention_mode, 'nearest')
+            self.attention = PositionAttention(
+                max_length=config.dataset_max_length + 1,  # additional stop token
+                mode=mode,
+            )
+        elif config.model_vision_attention == 'attention':
+            self.attention = Attention(
+                max_length=config.dataset_max_length + 1,  # additional stop token
+                n_feature=8*32,
+            )
+        else:
+            raise Exception(f'{config.model_vision_attention} is not valid.')
+        self.cls = nn.Linear(self.out_channels, self.charset.num_classes)
+
+        if config.model_vision_checkpoint is not None:
+            logging.info(f'Read vision model from {config.model_vision_checkpoint}.')
+            self.load(config.model_vision_checkpoint)
+
+    def forward(self, images, *args):
+        features = self.backbone(images)  # (N, E, H, W)
+        attn_vecs, attn_scores = self.attention(features)  # (N, T, E), (N, T, H, W)
+        logits = self.cls(attn_vecs) # (N, T, C)
+        pt_lengths = self._get_length(logits)
+
+        return {'feature': attn_vecs, 'logits': logits, 'pt_lengths': pt_lengths,
+                'attn_scores': attn_scores, 'loss_weight':self.loss_weight, 'name': 'vision'}

modules/resnet.py

@@ -0,0 +1,104 @@
+import math
+
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.model_zoo as model_zoo
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    "3x3 convolution with padding"
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv1x1(inplanes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers):
+        self.inplanes = 32
+        super(ResNet, self).__init__()
+        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(32)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.layer1 = self._make_layer(block, 32, layers[0], stride=2)
+        self.layer2 = self._make_layer(block, 64, layers[1], stride=1)
+        self.layer3 = self._make_layer(block, 128, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 256, layers[3], stride=1)
+        self.layer5 = self._make_layer(block, 512, layers[4], stride=1)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.layer5(x)
+        return x
+
+
+def resnet45():
+    return ResNet(BasicBlock, [3, 4, 6, 6, 3])

modules/transformer.py

@@ -0,0 +1,901 @@
+# pytorch 1.5.0
+import copy
+import math
+import warnings
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import Dropout, LayerNorm, Linear, Module, ModuleList, Parameter
+from torch.nn import functional as F
+from torch.nn.init import constant_, xavier_uniform_
+
+
+def multi_head_attention_forward(query,                           # type: Tensor
+                                 key,                             # type: Tensor
+                                 value,                           # type: Tensor
+                                 embed_dim_to_check,              # type: int
+                                 num_heads,                       # type: int
+                                 in_proj_weight,                  # type: Tensor
+                                 in_proj_bias,                    # type: Tensor
+                                 bias_k,                          # type: Optional[Tensor]
+                                 bias_v,                          # type: Optional[Tensor]
+                                 add_zero_attn,                   # type: bool
+                                 dropout_p,                       # type: float
+                                 out_proj_weight,                 # type: Tensor
+                                 out_proj_bias,                   # type: Tensor
+                                 training=True,                   # type: bool
+                                 key_padding_mask=None,           # type: Optional[Tensor]
+                                 need_weights=True,               # type: bool
+                                 attn_mask=None,                  # type: Optional[Tensor]
+                                 use_separate_proj_weight=False,  # type: bool
+                                 q_proj_weight=None,              # type: Optional[Tensor]
+                                 k_proj_weight=None,              # type: Optional[Tensor]
+                                 v_proj_weight=None,              # type: Optional[Tensor]
+                                 static_k=None,                   # type: Optional[Tensor]
+                                 static_v=None                    # type: Optional[Tensor]
+                                 ):
+    # type: (...) -> Tuple[Tensor, Optional[Tensor]]
+    r"""
+    Args:
+        query, key, value: map a query and a set of key-value pairs to an output.
+            See "Attention Is All You Need" for more details.
+        embed_dim_to_check: total dimension of the model.
+        num_heads: parallel attention heads.
+        in_proj_weight, in_proj_bias: input projection weight and bias.
+        bias_k, bias_v: bias of the key and value sequences to be added at dim=0.
+        add_zero_attn: add a new batch of zeros to the key and
+                       value sequences at dim=1.
+        dropout_p: probability of an element to be zeroed.
+        out_proj_weight, out_proj_bias: the output projection weight and bias.
+        training: apply dropout if is ``True``.
+        key_padding_mask: if provided, specified padding elements in the key will
+            be ignored by the attention. This is an binary mask. When the value is True,
+            the corresponding value on the attention layer will be filled with -inf.
+        need_weights: output attn_output_weights.
+        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
+            the batches while a 3D mask allows to specify a different mask for the entries of each batch.
+        use_separate_proj_weight: the function accept the proj. weights for query, key,
+            and value in different forms. If false, in_proj_weight will be used, which is
+            a combination of q_proj_weight, k_proj_weight, v_proj_weight.
+        q_proj_weight, k_proj_weight, v_proj_weight, in_proj_bias: input projection weight and bias.
+        static_k, static_v: static key and value used for attention operators.
+    Shape:
+        Inputs:
+        - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
+          If a ByteTensor is provided, the non-zero positions will be ignored while the zero positions
+          will be unchanged. If a BoolTensor is provided, the positions with the
+          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
+        - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
+          3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
+          S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked
+          positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
+          while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
+          are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
+          is provided, it will be added to the attention weight.
+        - static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
+          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
+        - static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
+          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
+        Outputs:
+        - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
+          E is the embedding dimension.
+        - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
+          L is the target sequence length, S is the source sequence length.
+    """
+    # if not torch.jit.is_scripting():
+    #     tens_ops = (query, key, value, in_proj_weight, in_proj_bias, bias_k, bias_v,
+    #                 out_proj_weight, out_proj_bias)
+    #     if any([type(t) is not Tensor for t in tens_ops]) and has_torch_function(tens_ops):
+    #         return handle_torch_function(
+    #             multi_head_attention_forward, tens_ops, query, key, value,
+    #             embed_dim_to_check, num_heads, in_proj_weight, in_proj_bias,
+    #             bias_k, bias_v, add_zero_attn, dropout_p, out_proj_weight,
+    #             out_proj_bias, training=training, key_padding_mask=key_padding_mask,
+    #             need_weights=need_weights, attn_mask=attn_mask,
+    #             use_separate_proj_weight=use_separate_proj_weight,
+    #             q_proj_weight=q_proj_weight, k_proj_weight=k_proj_weight,
+    #             v_proj_weight=v_proj_weight, static_k=static_k, static_v=static_v)
+    tgt_len, bsz, embed_dim = query.size()
+    assert embed_dim == embed_dim_to_check
+    assert key.size() == value.size()
+
+    head_dim = embed_dim // num_heads
+    assert head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"
+    scaling = float(head_dim) ** -0.5
+
+    if not use_separate_proj_weight:
+        if torch.equal(query, key) and torch.equal(key, value):
+            # self-attention
+            q, k, v = F.linear(query, in_proj_weight, in_proj_bias).chunk(3, dim=-1)
+
+        elif torch.equal(key, value):
+            # encoder-decoder attention
+            # This is inline in_proj function with in_proj_weight and in_proj_bias
+            _b = in_proj_bias
+            _start = 0
+            _end = embed_dim
+            _w = in_proj_weight[_start:_end, :]
+            if _b is not None:
+                _b = _b[_start:_end]
+            q = F.linear(query, _w, _b)
+
+            if key is None:
+                assert value is None
+                k = None
+                v = None
+            else:
+
+                # This is inline in_proj function with in_proj_weight and in_proj_bias
+                _b = in_proj_bias
+                _start = embed_dim
+                _end = None
+                _w = in_proj_weight[_start:, :]
+                if _b is not None:
+                    _b = _b[_start:]
+                k, v = F.linear(key, _w, _b).chunk(2, dim=-1)
+
+        else:
+            # This is inline in_proj function with in_proj_weight and in_proj_bias
+            _b = in_proj_bias
+            _start = 0
+            _end = embed_dim
+            _w = in_proj_weight[_start:_end, :]
+            if _b is not None:
+                _b = _b[_start:_end]
+            q = F.linear(query, _w, _b)
+
+            # This is inline in_proj function with in_proj_weight and in_proj_bias
+            _b = in_proj_bias
+            _start = embed_dim
+            _end = embed_dim * 2
+            _w = in_proj_weight[_start:_end, :]
+            if _b is not None:
+                _b = _b[_start:_end]
+            k = F.linear(key, _w, _b)
+
+            # This is inline in_proj function with in_proj_weight and in_proj_bias
+            _b = in_proj_bias
+            _start = embed_dim * 2
+            _end = None
+            _w = in_proj_weight[_start:, :]
+            if _b is not None:
+                _b = _b[_start:]
+            v = F.linear(value, _w, _b)
+    else:
+        q_proj_weight_non_opt = torch.jit._unwrap_optional(q_proj_weight)
+        len1, len2 = q_proj_weight_non_opt.size()
+        assert len1 == embed_dim and len2 == query.size(-1)
+
+        k_proj_weight_non_opt = torch.jit._unwrap_optional(k_proj_weight)
+        len1, len2 = k_proj_weight_non_opt.size()
+        assert len1 == embed_dim and len2 == key.size(-1)
+
+        v_proj_weight_non_opt = torch.jit._unwrap_optional(v_proj_weight)
+        len1, len2 = v_proj_weight_non_opt.size()
+        assert len1 == embed_dim and len2 == value.size(-1)
+
+        if in_proj_bias is not None:
+            q = F.linear(query, q_proj_weight_non_opt, in_proj_bias[0:embed_dim])
+            k = F.linear(key, k_proj_weight_non_opt, in_proj_bias[embed_dim:(embed_dim * 2)])
+            v = F.linear(value, v_proj_weight_non_opt, in_proj_bias[(embed_dim * 2):])
+        else:
+            q = F.linear(query, q_proj_weight_non_opt, in_proj_bias)
+            k = F.linear(key, k_proj_weight_non_opt, in_proj_bias)
+            v = F.linear(value, v_proj_weight_non_opt, in_proj_bias)
+    q = q * scaling
+
+    if attn_mask is not None:
+        assert attn_mask.dtype == torch.float32 or attn_mask.dtype == torch.float64 or \
+            attn_mask.dtype == torch.float16 or attn_mask.dtype == torch.uint8 or attn_mask.dtype == torch.bool, \
+            'Only float, byte, and bool types are supported for attn_mask, not {}'.format(attn_mask.dtype)
+        if attn_mask.dtype == torch.uint8:
+            warnings.warn("Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
+            attn_mask = attn_mask.to(torch.bool)
+
+        if attn_mask.dim() == 2:
+            attn_mask = attn_mask.unsqueeze(0)
+            if list(attn_mask.size()) != [1, query.size(0), key.size(0)]:
+                raise RuntimeError('The size of the 2D attn_mask is not correct.')
+        elif attn_mask.dim() == 3:
+            if list(attn_mask.size()) != [bsz * num_heads, query.size(0), key.size(0)]:
+                raise RuntimeError('The size of the 3D attn_mask is not correct.')
+        else:
+            raise RuntimeError("attn_mask's dimension {} is not supported".format(attn_mask.dim()))
+        # attn_mask's dim is 3 now.
+
+    # # convert ByteTensor key_padding_mask to bool
+    # if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:
+    #     warnings.warn("Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
+    #     key_padding_mask = key_padding_mask.to(torch.bool)
+
+    if bias_k is not None and bias_v is not None:
+        if static_k is None and static_v is None:
+            k = torch.cat([k, bias_k.repeat(1, bsz, 1)])
+            v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
+            if attn_mask is not None:
+                attn_mask = pad(attn_mask, (0, 1))
+            if key_padding_mask is not None:
+                key_padding_mask = pad(key_padding_mask, (0, 1))
+        else:
+            assert static_k is None, "bias cannot be added to static key."
+            assert static_v is None, "bias cannot be added to static value."
+    else:
+        assert bias_k is None
+        assert bias_v is None
+
+    q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
+    if k is not None:
+        k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
+    if v is not None:
+        v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
+
+    if static_k is not None:
+        assert static_k.size(0) == bsz * num_heads
+        assert static_k.size(2) == head_dim
+        k = static_k
+
+    if static_v is not None:
+        assert static_v.size(0) == bsz * num_heads
+        assert static_v.size(2) == head_dim
+        v = static_v
+
+    src_len = k.size(1)
+
+    if key_padding_mask is not None:
+        assert key_padding_mask.size(0) == bsz
+        assert key_padding_mask.size(1) == src_len
+
+    if add_zero_attn:
+        src_len += 1
+        k = torch.cat([k, torch.zeros((k.size(0), 1) + k.size()[2:], dtype=k.dtype, device=k.device)], dim=1)
+        v = torch.cat([v, torch.zeros((v.size(0), 1) + v.size()[2:], dtype=v.dtype, device=v.device)], dim=1)
+        if attn_mask is not None:
+            attn_mask = pad(attn_mask, (0, 1))
+        if key_padding_mask is not None:
+            key_padding_mask = pad(key_padding_mask, (0, 1))
+
+    attn_output_weights = torch.bmm(q, k.transpose(1, 2))
+    assert list(attn_output_weights.size()) == [bsz * num_heads, tgt_len, src_len]
+
+    if attn_mask is not None:
+        if attn_mask.dtype == torch.bool:
+            attn_output_weights.masked_fill_(attn_mask, float('-inf'))
+        else:
+            attn_output_weights += attn_mask
+
+
+    if key_padding_mask is not None:
+        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
+        attn_output_weights = attn_output_weights.masked_fill(
+            key_padding_mask.unsqueeze(1).unsqueeze(2),
+            float('-inf'),
+        )
+        attn_output_weights = attn_output_weights.view(bsz * num_heads, tgt_len, src_len)
+
+    attn_output_weights = F.softmax(
+        attn_output_weights, dim=-1)
+    attn_output_weights = F.dropout(attn_output_weights, p=dropout_p, training=training)
+
+    attn_output = torch.bmm(attn_output_weights, v)
+    assert list(attn_output.size()) == [bsz * num_heads, tgt_len, head_dim]
+    attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
+    attn_output = F.linear(attn_output, out_proj_weight, out_proj_bias)
+
+    if need_weights:
+        # average attention weights over heads
+        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
+        return attn_output, attn_output_weights.sum(dim=1) / num_heads
+    else:
+        return attn_output, None
+
+class MultiheadAttention(Module):
+    r"""Allows the model to jointly attend to information
+    from different representation subspaces.
+    See reference: Attention Is All You Need
+    .. math::
+        \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
+        \text{where} head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)
+    Args:
+        embed_dim: total dimension of the model.
+        num_heads: parallel attention heads.
+        dropout: a Dropout layer on attn_output_weights. Default: 0.0.
+        bias: add bias as module parameter. Default: True.
+        add_bias_kv: add bias to the key and value sequences at dim=0.
+        add_zero_attn: add a new batch of zeros to the key and
+                       value sequences at dim=1.
+        kdim: total number of features in key. Default: None.
+        vdim: total number of features in value. Default: None.
+        Note: if kdim and vdim are None, they will be set to embed_dim such that
+        query, key, and value have the same number of features.
+    Examples::
+        >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
+        >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
+    """
+    # __annotations__ = {
+    #     'bias_k': torch._jit_internal.Optional[torch.Tensor],
+    #     'bias_v': torch._jit_internal.Optional[torch.Tensor],
+    # }
+    __constants__ = ['q_proj_weight', 'k_proj_weight', 'v_proj_weight', 'in_proj_weight']
+
+    def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None):
+        super(MultiheadAttention, self).__init__()
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim
+
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
+
+        if self._qkv_same_embed_dim is False:
+            self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))
+            self.k_proj_weight = Parameter(torch.Tensor(embed_dim, self.kdim))
+            self.v_proj_weight = Parameter(torch.Tensor(embed_dim, self.vdim))
+            self.register_parameter('in_proj_weight', None)
+        else:
+            self.in_proj_weight = Parameter(torch.empty(3 * embed_dim, embed_dim))
+            self.register_parameter('q_proj_weight', None)
+            self.register_parameter('k_proj_weight', None)
+            self.register_parameter('v_proj_weight', None)
+
+        if bias:
+            self.in_proj_bias = Parameter(torch.empty(3 * embed_dim))
+        else:
+            self.register_parameter('in_proj_bias', None)
+        self.out_proj = Linear(embed_dim, embed_dim, bias=bias)
+
+        if add_bias_kv:
+            self.bias_k = Parameter(torch.empty(1, 1, embed_dim))
+            self.bias_v = Parameter(torch.empty(1, 1, embed_dim))
+        else:
+            self.bias_k = self.bias_v = None
+
+        self.add_zero_attn = add_zero_attn
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        if self._qkv_same_embed_dim:
+            xavier_uniform_(self.in_proj_weight)
+        else:
+            xavier_uniform_(self.q_proj_weight)
+            xavier_uniform_(self.k_proj_weight)
+            xavier_uniform_(self.v_proj_weight)
+
+        if self.in_proj_bias is not None:
+            constant_(self.in_proj_bias, 0.)
+            constant_(self.out_proj.bias, 0.)
+        if self.bias_k is not None:
+            xavier_normal_(self.bias_k)
+        if self.bias_v is not None:
+            xavier_normal_(self.bias_v)
+
+    def __setstate__(self, state):
+        # Support loading old MultiheadAttention checkpoints generated by v1.1.0
+        if '_qkv_same_embed_dim' not in state:
+            state['_qkv_same_embed_dim'] = True
+
+        super(MultiheadAttention, self).__setstate__(state)
+
+    def forward(self, query, key, value, key_padding_mask=None,
+                need_weights=True, attn_mask=None):
+        # type: (Tensor, Tensor, Tensor, Optional[Tensor], bool, Optional[Tensor]) -> Tuple[Tensor, Optional[Tensor]]
+        r"""
+    Args:
+        query, key, value: map a query and a set of key-value pairs to an output.
+            See "Attention Is All You Need" for more details.
+        key_padding_mask: if provided, specified padding elements in the key will
+            be ignored by the attention. This is an binary mask. When the value is True,
+            the corresponding value on the attention layer will be filled with -inf.
+        need_weights: output attn_output_weights.
+        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
+            the batches while a 3D mask allows to specify a different mask for the entries of each batch.
+    Shape:
+        - Inputs:
+        - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
+          If a ByteTensor is provided, the non-zero positions will be ignored while the position
+          with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the
+          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
+        - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
+          3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
+          S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked
+          positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
+          while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
+          is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
+          is provided, it will be added to the attention weight.
+        - Outputs:
+        - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
+          E is the embedding dimension.
+        - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
+          L is the target sequence length, S is the source sequence length.
+        """
+        if not self._qkv_same_embed_dim:
+            return multi_head_attention_forward(
+                query, key, value, self.embed_dim, self.num_heads,
+                self.in_proj_weight, self.in_proj_bias,
+                self.bias_k, self.bias_v, self.add_zero_attn,
+                self.dropout, self.out_proj.weight, self.out_proj.bias,
+                training=self.training,
+                key_padding_mask=key_padding_mask, need_weights=need_weights,
+                attn_mask=attn_mask, use_separate_proj_weight=True,
+                q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,
+                v_proj_weight=self.v_proj_weight)
+        else:
+            return multi_head_attention_forward(
+                query, key, value, self.embed_dim, self.num_heads,
+                self.in_proj_weight, self.in_proj_bias,
+                self.bias_k, self.bias_v, self.add_zero_attn,
+                self.dropout, self.out_proj.weight, self.out_proj.bias,
+                training=self.training,
+                key_padding_mask=key_padding_mask, need_weights=need_weights,
+                attn_mask=attn_mask)
+
+
+class Transformer(Module):
+    r"""A transformer model. User is able to modify the attributes as needed. The architecture
+    is based on the paper "Attention Is All You Need". Ashish Vaswani, Noam Shazeer,
+    Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and
+    Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information
+    Processing Systems, pages 6000-6010. Users can build the BERT(https://arxiv.org/abs/1810.04805)
+    model with corresponding parameters.
+
+    Args:
+        d_model: the number of expected features in the encoder/decoder inputs (default=512).
+        nhead: the number of heads in the multiheadattention models (default=8).
+        num_encoder_layers: the number of sub-encoder-layers in the encoder (default=6).
+        num_decoder_layers: the number of sub-decoder-layers in the decoder (default=6).
+        dim_feedforward: the dimension of the feedforward network model (default=2048).
+        dropout: the dropout value (default=0.1).
+        activation: the activation function of encoder/decoder intermediate layer, relu or gelu (default=relu).
+        custom_encoder: custom encoder (default=None).
+        custom_decoder: custom decoder (default=None).
+
+    Examples::
+        >>> transformer_model = nn.Transformer(nhead=16, num_encoder_layers=12)
+        >>> src = torch.rand((10, 32, 512))
+        >>> tgt = torch.rand((20, 32, 512))
+        >>> out = transformer_model(src, tgt)
+
+    Note: A full example to apply nn.Transformer module for the word language model is available in
+    https://github.com/pytorch/examples/tree/master/word_language_model
+    """
+
+    def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
+                 num_decoder_layers=6, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", custom_encoder=None, custom_decoder=None):
+        super(Transformer, self).__init__()
+
+        if custom_encoder is not None:
+            self.encoder = custom_encoder
+        else:
+            encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, activation)
+            encoder_norm = LayerNorm(d_model)
+            self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
+
+        if custom_decoder is not None:
+            self.decoder = custom_decoder
+        else:
+            decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout, activation)
+            decoder_norm = LayerNorm(d_model)
+            self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)
+
+        self._reset_parameters()
+
+        self.d_model = d_model
+        self.nhead = nhead
+
+    def forward(self, src, tgt, src_mask=None, tgt_mask=None,
+                memory_mask=None, src_key_padding_mask=None,
+                tgt_key_padding_mask=None, memory_key_padding_mask=None):
+        # type: (Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor]) -> Tensor  # noqa
+        r"""Take in and process masked source/target sequences.
+
+        Args:
+            src: the sequence to the encoder (required).
+            tgt: the sequence to the decoder (required).
+            src_mask: the additive mask for the src sequence (optional).
+            tgt_mask: the additive mask for the tgt sequence (optional).
+            memory_mask: the additive mask for the encoder output (optional).
+            src_key_padding_mask: the ByteTensor mask for src keys per batch (optional).
+            tgt_key_padding_mask: the ByteTensor mask for tgt keys per batch (optional).
+            memory_key_padding_mask: the ByteTensor mask for memory keys per batch (optional).
+
+        Shape:
+            - src: :math:`(S, N, E)`.
+            - tgt: :math:`(T, N, E)`.
+            - src_mask: :math:`(S, S)`.
+            - tgt_mask: :math:`(T, T)`.
+            - memory_mask: :math:`(T, S)`.
+            - src_key_padding_mask: :math:`(N, S)`.
+            - tgt_key_padding_mask: :math:`(N, T)`.
+            - memory_key_padding_mask: :math:`(N, S)`.
+
+            Note: [src/tgt/memory]_mask ensures that position i is allowed to attend the unmasked
+            positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
+            while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
+            are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
+            is provided, it will be added to the attention weight. 
+            [src/tgt/memory]_key_padding_mask provides specified elements in the key to be ignored by
+            the attention. If a ByteTensor is provided, the non-zero positions will be ignored while the zero
+            positions will be unchanged. If a BoolTensor is provided, the positions with the
+            value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
+
+            - output: :math:`(T, N, E)`.
+
+            Note: Due to the multi-head attention architecture in the transformer model,
+            the output sequence length of a transformer is same as the input sequence
+            (i.e. target) length of the decode.
+
+            where S is the source sequence length, T is the target sequence length, N is the
+            batch size, E is the feature number
+
+        Examples:
+            >>> output = transformer_model(src, tgt, src_mask=src_mask, tgt_mask=tgt_mask)
+        """
+
+        if src.size(1) != tgt.size(1):
+            raise RuntimeError("the batch number of src and tgt must be equal")
+
+        if src.size(2) != self.d_model or tgt.size(2) != self.d_model:
+            raise RuntimeError("the feature number of src and tgt must be equal to d_model")
+
+        memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask)
+        output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
+                              tgt_key_padding_mask=tgt_key_padding_mask,
+                              memory_key_padding_mask=memory_key_padding_mask)
+        return output
+
+    def generate_square_subsequent_mask(self, sz):
+        r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf').
+            Unmasked positions are filled with float(0.0).
+        """
+        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
+        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+        return mask
+
+    def _reset_parameters(self):
+        r"""Initiate parameters in the transformer model."""
+
+        for p in self.parameters():
+            if p.dim() > 1:
+                xavier_uniform_(p)
+
+
+class TransformerEncoder(Module):
+    r"""TransformerEncoder is a stack of N encoder layers
+
+    Args:
+        encoder_layer: an instance of the TransformerEncoderLayer() class (required).
+        num_layers: the number of sub-encoder-layers in the encoder (required).
+        norm: the layer normalization component (optional).
+
+    Examples::
+        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
+        >>> transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
+        >>> src = torch.rand(10, 32, 512)
+        >>> out = transformer_encoder(src)
+    """
+    __constants__ = ['norm']
+
+    def __init__(self, encoder_layer, num_layers, norm=None):
+        super(TransformerEncoder, self).__init__()
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+
+    def forward(self, src, mask=None, src_key_padding_mask=None):
+        # type: (Tensor, Optional[Tensor], Optional[Tensor]) -> Tensor
+        r"""Pass the input through the encoder layers in turn.
+
+        Args:
+            src: the sequence to the encoder (required).
+            mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        output = src
+
+        for i, mod in enumerate(self.layers):
+            output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        return output
+
+
+class TransformerDecoder(Module):
+    r"""TransformerDecoder is a stack of N decoder layers
+
+    Args:
+        decoder_layer: an instance of the TransformerDecoderLayer() class (required).
+        num_layers: the number of sub-decoder-layers in the decoder (required).
+        norm: the layer normalization component (optional).
+
+    Examples::
+        >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
+        >>> transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
+        >>> memory = torch.rand(10, 32, 512)
+        >>> tgt = torch.rand(20, 32, 512)
+        >>> out = transformer_decoder(tgt, memory)
+    """
+    __constants__ = ['norm']
+
+    def __init__(self, decoder_layer, num_layers, norm=None):
+        super(TransformerDecoder, self).__init__()
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+
+    def forward(self, tgt, memory, memory2=None, tgt_mask=None,
+                memory_mask=None, memory_mask2=None, tgt_key_padding_mask=None,
+                memory_key_padding_mask=None, memory_key_padding_mask2=None):
+        # type: (Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor]) -> Tensor
+        r"""Pass the inputs (and mask) through the decoder layer in turn.
+
+        Args:
+            tgt: the sequence to the decoder (required).
+            memory: the sequence from the last layer of the encoder (required).
+            tgt_mask: the mask for the tgt sequence (optional).
+            memory_mask: the mask for the memory sequence (optional).
+            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
+            memory_key_padding_mask: the mask for the memory keys per batch (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        output = tgt
+
+        for mod in self.layers:
+            output = mod(output, memory, memory2=memory2, tgt_mask=tgt_mask,
+                         memory_mask=memory_mask, memory_mask2=memory_mask2,
+                         tgt_key_padding_mask=tgt_key_padding_mask,
+                         memory_key_padding_mask=memory_key_padding_mask,
+                         memory_key_padding_mask2=memory_key_padding_mask2)
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        return output
+
+class TransformerEncoderLayer(Module):
+    r"""TransformerEncoderLayer is made up of self-attn and feedforward network.
+    This standard encoder layer is based on the paper "Attention Is All You Need".
+    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
+    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
+    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
+    in a different way during application.
+
+    Args:
+        d_model: the number of expected features in the input (required).
+        nhead: the number of heads in the multiheadattention models (required).
+        dim_feedforward: the dimension of the feedforward network model (default=2048).
+        dropout: the dropout value (default=0.1).
+        activation: the activation function of intermediate layer, relu or gelu (default=relu).
+
+    Examples::
+        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
+        >>> src = torch.rand(10, 32, 512)
+        >>> out = encoder_layer(src)
+    """
+
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, 
+                 activation="relu", debug=False):
+        super(TransformerEncoderLayer, self).__init__()
+        self.debug = debug
+        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
+        # Implementation of Feedforward model
+        self.linear1 = Linear(d_model, dim_feedforward)
+        self.dropout = Dropout(dropout)
+        self.linear2 = Linear(dim_feedforward, d_model)
+
+        self.norm1 = LayerNorm(d_model)
+        self.norm2 = LayerNorm(d_model)
+        self.dropout1 = Dropout(dropout)
+        self.dropout2 = Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+
+    def __setstate__(self, state):
+        if 'activation' not in state:
+            state['activation'] = F.relu
+        super(TransformerEncoderLayer, self).__setstate__(state)
+
+    def forward(self, src, src_mask=None, src_key_padding_mask=None):
+        # type: (Tensor, Optional[Tensor], Optional[Tensor]) -> Tensor
+        r"""Pass the input through the encoder layer.
+
+        Args:
+            src: the sequence to the encoder layer (required).
+            src_mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        src2, attn = self.self_attn(src, src, src, attn_mask=src_mask,
+                              key_padding_mask=src_key_padding_mask)
+        if self.debug: self.attn = attn
+        src = src + self.dropout1(src2)
+        src = self.norm1(src)
+        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
+        src = src + self.dropout2(src2)
+        src = self.norm2(src)
+        
+        return src
+
+
+class TransformerDecoderLayer(Module):
+    r"""TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network.
+    This standard decoder layer is based on the paper "Attention Is All You Need".
+    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
+    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
+    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
+    in a different way during application.
+
+    Args:
+        d_model: the number of expected features in the input (required).
+        nhead: the number of heads in the multiheadattention models (required).
+        dim_feedforward: the dimension of the feedforward network model (default=2048).
+        dropout: the dropout value (default=0.1).
+        activation: the activation function of intermediate layer, relu or gelu (default=relu).
+
+    Examples::
+        >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
+        >>> memory = torch.rand(10, 32, 512)
+        >>> tgt = torch.rand(20, 32, 512)
+        >>> out = decoder_layer(tgt, memory)
+    """
+
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, 
+                 activation="relu", self_attn=True, siamese=False, debug=False):
+        super(TransformerDecoderLayer, self).__init__()
+        self.has_self_attn, self.siamese = self_attn, siamese
+        self.debug = debug
+        if self.has_self_attn:
+            self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
+            self.norm1 = LayerNorm(d_model)
+            self.dropout1 = Dropout(dropout)
+        self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
+        # Implementation of Feedforward model
+        self.linear1 = Linear(d_model, dim_feedforward)
+        self.dropout = Dropout(dropout)
+        self.linear2 = Linear(dim_feedforward, d_model)
+
+        self.norm2 = LayerNorm(d_model)
+        self.norm3 = LayerNorm(d_model)
+        self.dropout2 = Dropout(dropout)
+        self.dropout3 = Dropout(dropout)
+        if self.siamese:
+            self.multihead_attn2 = MultiheadAttention(d_model, nhead, dropout=dropout)
+
+        self.activation = _get_activation_fn(activation)
+
+    def __setstate__(self, state):
+        if 'activation' not in state:
+            state['activation'] = F.relu
+        super(TransformerDecoderLayer, self).__setstate__(state)
+
+    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None,
+                tgt_key_padding_mask=None, memory_key_padding_mask=None,
+                memory2=None, memory_mask2=None, memory_key_padding_mask2=None):
+        # type: (Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor]) -> Tensor
+        r"""Pass the inputs (and mask) through the decoder layer.
+
+        Args:
+            tgt: the sequence to the decoder layer (required).
+            memory: the sequence from the last layer of the encoder (required).
+            tgt_mask: the mask for the tgt sequence (optional).
+            memory_mask: the mask for the memory sequence (optional).
+            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
+            memory_key_padding_mask: the mask for the memory keys per batch (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        if self.has_self_attn:
+            tgt2, attn = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,
+                                key_padding_mask=tgt_key_padding_mask)
+            tgt = tgt + self.dropout1(tgt2)
+            tgt = self.norm1(tgt)
+            if self.debug: self.attn = attn
+        tgt2, attn2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask,
+                                   key_padding_mask=memory_key_padding_mask)
+        if self.debug: self.attn2 = attn2
+
+        if self.siamese:
+            tgt3, attn3 = self.multihead_attn2(tgt, memory2, memory2, attn_mask=memory_mask2,
+                            key_padding_mask=memory_key_padding_mask2)
+            tgt = tgt + self.dropout2(tgt3)
+            if self.debug: self.attn3 = attn3
+
+        tgt = tgt + self.dropout2(tgt2)
+        tgt = self.norm2(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout3(tgt2)
+        tgt = self.norm3(tgt)
+        
+        return tgt
+
+
+def _get_clones(module, N):
+    return ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def _get_activation_fn(activation):
+    if activation == "relu":
+        return F.relu
+    elif activation == "gelu":
+        return F.gelu
+
+    raise RuntimeError("activation should be relu/gelu, not {}".format(activation))
+
+
+class PositionalEncoding(nn.Module):
+    r"""Inject some information about the relative or absolute position of the tokens
+        in the sequence. The positional encodings have the same dimension as
+        the embeddings, so that the two can be summed. Here, we use sine and cosine
+        functions of different frequencies.
+    .. math::
+        \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
+        \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
+        \text{where pos is the word position and i is the embed idx)
+    Args:
+        d_model: the embed dim (required).
+        dropout: the dropout value (default=0.1).
+        max_len: the max. length of the incoming sequence (default=5000).
+    Examples:
+        >>> pos_encoder = PositionalEncoding(d_model)
+    """
+
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        r"""Inputs of forward function
+        Args:
+            x: the sequence fed to the positional encoder model (required).
+        Shape:
+            x: [sequence length, batch size, embed dim]
+            output: [sequence length, batch size, embed dim]
+        Examples:
+            >>> output = pos_encoder(x)
+        """
+
+        x = x + self.pe[:x.size(0), :]
+        return self.dropout(x)
+
+
+if __name__ == '__main__':
+    transformer_model = Transformer(nhead=16, num_encoder_layers=12)
+    src = torch.rand((10, 32, 512))
+    tgt = torch.rand((20, 32, 512))
+    out = transformer_model(src, tgt)
+    print(out)

notebooks/dataset-text.ipynb

@@ -0,0 +1,159 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.chdir('..')\n",
+    "from dataset import *\n",
+    "torch.set_printoptions(sci_mode=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Construct dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data = TextDataset('data/Vocabulary_train_v2.csv', is_training=False, smooth_label=True, smooth_factor=0.1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data = DataBunch.create(train_ds=data, valid_ds=None, bs=6)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "x, y = data.one_batch(); x, y"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "x[0].shape, x[1].shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "y[0].shape, y[1].shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "x[0].argmax(-1) - y[0].argmax(-1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "x[0].argmax(-1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "y[0].argmax(-1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "x[0][0,0]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# test SpellingMutation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "probs = {'pn0': 0., 'pn1': 0., 'pn2': 0., 'pt0': 1.0, 'pt1': 1.0}\n",
+    "charset = CharsetMapper('data/charset_36.txt')\n",
+    "sm = SpellingMutation(charset=charset, **probs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sm('*a-aa')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

notebooks/dataset.ipynb

@@ -0,0 +1,298 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.chdir('..')\n",
+    "from dataset import *"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "import logging\n",
+    "from torchvision.transforms import ToPILImage\n",
+    "from torchvision.utils import make_grid\n",
+    "from IPython.display import display\n",
+    "from torch.utils.data import ConcatDataset\n",
+    "charset = CharsetMapper('data/charset_36.txt')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def show_all(dl, iter_size=None):\n",
+    "    if iter_size is None: iter_size = len(dl)\n",
+    "    for i, item in enumerate(dl):\n",
+    "        if i >= iter_size:\n",
+    "             break\n",
+    "        image = item[0]\n",
+    "        label = item[1][0]\n",
+    "        length = item[1][1]\n",
+    "        print(f'iter {i}:', [charset.get_text(label[j][0: length[j]].argmax(-1), padding=False) for j in range(bs)])\n",
+    "        display(ToPILImage()(make_grid(item[0].cpu())))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Construct dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data1 = ImageDataset('data/training/ST', is_training=True);data1  # is_training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "bs=64\n",
+    "data2 = ImageDataBunch.create(train_ds=data1, valid_ds=None, bs=bs, num_workers=1);data2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#data3 = data2.normalize(imagenet_stats);data3\n",
+    "data3 = data2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "show_all(data3.train_dl, 4)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Add dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "kwargs = {'data_aug': False, 'is_training': False}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data1 = ImageDataset('data/evaluation/IIIT5k_3000', **kwargs);data1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data2 = ImageDataset('data/evaluation/SVT', **kwargs);data2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data3 = ConcatDataset([data1, data2])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "bs=64\n",
+    "data4 = ImageDataBunch.create(train_ds=data1, valid_ds=data3, bs=bs, num_workers=1);data4"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "len(data4.train_dl), len(data4.valid_dl)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "show_all(data4.train_dl, 4)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# TEST"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "len(data4.valid_dl)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import time\n",
+    "niter = 1000\n",
+    "start = time.time()\n",
+    "for i, item in enumerate(progress_bar(data4.valid_dl)):\n",
+    "    if i % niter == 0 and i > 0:\n",
+    "        print(i, (time.time() - start) / niter)\n",
+    "        start = time.time()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "num = 20\n",
+    "index = 6\n",
+    "plt.figure(figsize=(20, 10))\n",
+    "for i in range(num):\n",
+    "    plt.subplot(num // 4, 4, i+1)\n",
+    "    plt.imshow(data4.train_ds[i][0].data.numpy().transpose(1,2,0))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def show(path, image_key):\n",
+    "    with lmdb.open(str(path), readonly=True, lock=False, readahead=False, meminit=False).begin(write=False) as txn:\n",
+    "        imgbuf = txn.get(image_key.encode())  # image\n",
+    "        buf = six.BytesIO()\n",
+    "        buf.write(imgbuf)\n",
+    "        buf.seek(0)\n",
+    "        with warnings.catch_warnings():\n",
+    "            warnings.simplefilter(\"ignore\", UserWarning) # EXIF warning from TiffPlugin\n",
+    "            x = PIL.Image.open(buf).convert('RGB')\n",
+    "        print(x.size)\n",
+    "        plt.imshow(x)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image_key = 'image-003118258'\n",
+    "image_key = 'image-002780217'\n",
+    "image_key = 'image-002780218'\n",
+    "path = 'data/CVPR2016'\n",
+    "show(path, image_key)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image_key = 'image-004668347'\n",
+    "image_key = 'image-006128516'\n",
+    "path = 'data/NIPS2014'\n",
+    "show(path, image_key)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image_key = 'image-004668347'\n",
+    "image_key = 'image-000002420'\n",
+    "path = 'data/IIIT5K_3000'\n",
+    "show(path, image_key)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

notebooks/prepare_wikitext103.ipynb

@@ -0,0 +1,468 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 82841986 is_char and is_digit"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 82075350 regrex non-ascii and none-digit"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 86460763 left"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import random\n",
+    "import re\n",
+    "import pandas as pd"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "max_length = 25\n",
+    "min_length = 1\n",
+    "root = '../data'\n",
+    "charset = 'abcdefghijklmnopqrstuvwxyz'\n",
+    "digits = '0123456789'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def is_char(text, ratio=0.5):\n",
+    "    text = text.lower()\n",
+    "    length = max(len(text), 1)\n",
+    "    char_num = sum([t in charset for t in text])\n",
+    "    if char_num < min_length: return False\n",
+    "    if char_num / length < ratio: return False\n",
+    "    return True\n",
+    "\n",
+    "def is_digit(text, ratio=0.5):\n",
+    "    length = max(len(text), 1)\n",
+    "    digit_num = sum([t in digits for t in text])\n",
+    "    if digit_num / length < ratio: return False\n",
+    "    return True"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# generate training dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open('/tmp/wikitext-103/wiki.train.tokens', 'r') as file:\n",
+    "    lines = file.readlines()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "inp, gt = [], []\n",
+    "for line in lines:\n",
+    "    token = line.lower().split()\n",
+    "    for text in token:\n",
+    "        text = re.sub('[^0-9a-zA-Z]+', '', text)\n",
+    "        if len(text) < min_length:\n",
+    "            # print('short-text', text)\n",
+    "            continue\n",
+    "        if len(text) > max_length:\n",
+    "            # print('long-text', text)\n",
+    "            continue\n",
+    "        inp.append(text)\n",
+    "        gt.append(text)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_voc = os.path.join(root, 'WikiText-103.csv')\n",
+    "pd.DataFrame({'inp':inp, 'gt':gt}).to_csv(train_voc, index=None, sep='\\t')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "86460763"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "len(inp)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['valkyria',\n",
+       " 'chronicles',\n",
+       " 'iii',\n",
+       " 'senj',\n",
+       " 'no',\n",
+       " 'valkyria',\n",
+       " '3',\n",
+       " 'unk',\n",
+       " 'chronicles',\n",
+       " 'japanese',\n",
+       " '3',\n",
+       " 'lit',\n",
+       " 'valkyria',\n",
+       " 'of',\n",
+       " 'the',\n",
+       " 'battlefield',\n",
+       " '3',\n",
+       " 'commonly',\n",
+       " 'referred',\n",
+       " 'to',\n",
+       " 'as',\n",
+       " 'valkyria',\n",
+       " 'chronicles',\n",
+       " 'iii',\n",
+       " 'outside',\n",
+       " 'japan',\n",
+       " 'is',\n",
+       " 'a',\n",
+       " 'tactical',\n",
+       " 'role',\n",
+       " 'playing',\n",
+       " 'video',\n",
+       " 'game',\n",
+       " 'developed',\n",
+       " 'by',\n",
+       " 'sega',\n",
+       " 'and',\n",
+       " 'mediavision',\n",
+       " 'for',\n",
+       " 'the',\n",
+       " 'playstation',\n",
+       " 'portable',\n",
+       " 'released',\n",
+       " 'in',\n",
+       " 'january',\n",
+       " '2011',\n",
+       " 'in',\n",
+       " 'japan',\n",
+       " 'it',\n",
+       " 'is',\n",
+       " 'the',\n",
+       " 'third',\n",
+       " 'game',\n",
+       " 'in',\n",
+       " 'the',\n",
+       " 'valkyria',\n",
+       " 'series',\n",
+       " 'employing',\n",
+       " 'the',\n",
+       " 'same',\n",
+       " 'fusion',\n",
+       " 'of',\n",
+       " 'tactical',\n",
+       " 'and',\n",
+       " 'real',\n",
+       " 'time',\n",
+       " 'gameplay',\n",
+       " 'as',\n",
+       " 'its',\n",
+       " 'predecessors',\n",
+       " 'the',\n",
+       " 'story',\n",
+       " 'runs',\n",
+       " 'parallel',\n",
+       " 'to',\n",
+       " 'the',\n",
+       " 'first',\n",
+       " 'game',\n",
+       " 'and',\n",
+       " 'follows',\n",
+       " 'the',\n",
+       " 'nameless',\n",
+       " 'a',\n",
+       " 'penal',\n",
+       " 'military',\n",
+       " 'unit',\n",
+       " 'serving',\n",
+       " 'the',\n",
+       " 'nation',\n",
+       " 'of',\n",
+       " 'gallia',\n",
+       " 'during',\n",
+       " 'the',\n",
+       " 'second',\n",
+       " 'europan',\n",
+       " 'war',\n",
+       " 'who',\n",
+       " 'perform',\n",
+       " 'secret',\n",
+       " 'black']"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "inp[:100]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# generate evaluation dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def disturb(word, degree, p=0.3):\n",
+    "    if len(word) // 2 < degree: return word\n",
+    "    if is_digit(word): return word\n",
+    "    if random.random() < p: return word\n",
+    "    else:\n",
+    "        index = list(range(len(word)))\n",
+    "        random.shuffle(index)\n",
+    "        index = index[:degree]\n",
+    "        new_word = []\n",
+    "        for i in range(len(word)):\n",
+    "            if i not in index: \n",
+    "                new_word.append(word[i])\n",
+    "                continue\n",
+    "            if (word[i] not in charset) and (word[i] not in digits):\n",
+    "                # special token\n",
+    "                new_word.append(word[i])\n",
+    "                continue\n",
+    "            op = random.random()\n",
+    "            if op < 0.1: # add\n",
+    "                new_word.append(random.choice(charset))\n",
+    "                new_word.append(word[i])\n",
+    "            elif op < 0.2: continue  # remove\n",
+    "            else: new_word.append(random.choice(charset))  # replace\n",
+    "        return ''.join(new_word)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "lines = inp\n",
+    "degree = 1\n",
+    "keep_num = 50000\n",
+    "\n",
+    "random.shuffle(lines)\n",
+    "part_lines = lines[:keep_num]\n",
+    "inp, gt = [], []\n",
+    "\n",
+    "for w in part_lines:\n",
+    "    w = w.strip().lower()\n",
+    "    new_w = disturb(w, degree)\n",
+    "    inp.append(new_w)\n",
+    "    gt.append(w)\n",
+    "    \n",
+    "eval_voc = os.path.join(root, f'WikiText-103_eval_d{degree}.csv')\n",
+    "pd.DataFrame({'inp':inp, 'gt':gt}).to_csv(eval_voc, index=None, sep='\\t')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[('high', 'high'),\n",
+       " ('vctoria', 'victoria'),\n",
+       " ('mains', 'mains'),\n",
+       " ('bi', 'by'),\n",
+       " ('13', '13'),\n",
+       " ('ticnet', 'ticket'),\n",
+       " ('basil', 'basic'),\n",
+       " ('cut', 'cut'),\n",
+       " ('aqarky', 'anarky'),\n",
+       " ('the', 'the'),\n",
+       " ('tqe', 'the'),\n",
+       " ('oc', 'of'),\n",
+       " ('diwpersal', 'dispersal'),\n",
+       " ('traffic', 'traffic'),\n",
+       " ('in', 'in'),\n",
+       " ('the', 'the'),\n",
+       " ('ti', 'to'),\n",
+       " ('professionalms', 'professionals'),\n",
+       " ('747', '747'),\n",
+       " ('in', 'in'),\n",
+       " ('and', 'and'),\n",
+       " ('exezutive', 'executive'),\n",
+       " ('n400', 'n400'),\n",
+       " ('yusic', 'music'),\n",
+       " ('s', 's'),\n",
+       " ('henri', 'henry'),\n",
+       " ('heard', 'heard'),\n",
+       " ('thousand', 'thousand'),\n",
+       " ('to', 'to'),\n",
+       " ('arhy', 'army'),\n",
+       " ('td', 'to'),\n",
+       " ('a', 'a'),\n",
+       " ('oall', 'hall'),\n",
+       " ('qind', 'kind'),\n",
+       " ('od', 'on'),\n",
+       " ('samfria', 'samaria'),\n",
+       " ('driveway', 'driveway'),\n",
+       " ('which', 'which'),\n",
+       " ('wotk', 'work'),\n",
+       " ('ak', 'as'),\n",
+       " ('persona', 'persona'),\n",
+       " ('s', 's'),\n",
+       " ('melbourne', 'melbourne'),\n",
+       " ('apong', 'along'),\n",
+       " ('fas', 'was'),\n",
+       " ('thea', 'then'),\n",
+       " ('permcy', 'percy'),\n",
+       " ('nnd', 'and'),\n",
+       " ('alan', 'alan'),\n",
+       " ('13', '13'),\n",
+       " ('matteos', 'matters'),\n",
+       " ('against', 'against'),\n",
+       " ('nefion', 'nexion'),\n",
+       " ('held', 'held'),\n",
+       " ('negative', 'negative'),\n",
+       " ('gogd', 'good'),\n",
+       " ('the', 'the'),\n",
+       " ('thd', 'the'),\n",
+       " ('groening', 'groening'),\n",
+       " ('tqe', 'the'),\n",
+       " ('cwould', 'would'),\n",
+       " ('fb', 'ft'),\n",
+       " ('uniten', 'united'),\n",
+       " ('kone', 'one'),\n",
+       " ('thiy', 'this'),\n",
+       " ('lanren', 'lauren'),\n",
+       " ('s', 's'),\n",
+       " ('thhe', 'the'),\n",
+       " ('is', 'is'),\n",
+       " ('modep', 'model'),\n",
+       " ('weird', 'weird'),\n",
+       " ('angwer', 'answer'),\n",
+       " ('imprisxnment', 'imprisonment'),\n",
+       " ('marpery', 'margery'),\n",
+       " ('eventuanly', 'eventually'),\n",
+       " ('in', 'in'),\n",
+       " ('donnoa', 'donna'),\n",
+       " ('ik', 'it'),\n",
+       " ('reached', 'reached'),\n",
+       " ('at', 'at'),\n",
+       " ('excxted', 'excited'),\n",
+       " ('ws', 'was'),\n",
+       " ('raes', 'rates'),\n",
+       " ('the', 'the'),\n",
+       " ('firsq', 'first'),\n",
+       " ('concluyed', 'concluded'),\n",
+       " ('recdorded', 'recorded'),\n",
+       " ('fhe', 'the'),\n",
+       " ('uegiment', 'regiment'),\n",
+       " ('a', 'a'),\n",
+       " ('glanes', 'planes'),\n",
+       " ('conyrol', 'control'),\n",
+       " ('thr', 'the'),\n",
+       " ('arrext', 'arrest'),\n",
+       " ('bth', 'both'),\n",
+       " ('forward', 'forward'),\n",
+       " ('allowdd', 'allowed'),\n",
+       " ('revealed', 'revealed'),\n",
+       " ('mayagement', 'management'),\n",
+       " ('normal', 'normal')]"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "list(zip(inp, gt))[:100]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

requirements.txt

@@ -0,0 +1,7 @@
+torch==1.1.0
+torchvision==0.3.0
+fastai==1.0.60
+LMDB
+Pillow
+opencv-python
+tensorboardX

tools/create_lmdb_dataset.py

@@ -0,0 +1,87 @@
+""" a modified version of CRNN torch repository https://github.com/bgshih/crnn/blob/master/tool/create_dataset.py """
+
+import fire
+import os
+import lmdb
+import cv2
+
+import numpy as np
+
+
+def checkImageIsValid(imageBin):
+    if imageBin is None:
+        return False
+    imageBuf = np.frombuffer(imageBin, dtype=np.uint8)
+    img = cv2.imdecode(imageBuf, cv2.IMREAD_GRAYSCALE)
+    imgH, imgW = img.shape[0], img.shape[1]
+    if imgH * imgW == 0:
+        return False
+    return True
+
+
+def writeCache(env, cache):
+    with env.begin(write=True) as txn:
+        for k, v in cache.items():
+            txn.put(k, v)
+
+
+def createDataset(inputPath, gtFile, outputPath, checkValid=True):
+    """
+    Create LMDB dataset for training and evaluation.
+    ARGS:
+        inputPath  : input folder path where starts imagePath
+        outputPath : LMDB output path
+        gtFile     : list of image path and label
+        checkValid : if true, check the validity of every image
+    """
+    os.makedirs(outputPath, exist_ok=True)
+    env = lmdb.open(outputPath, map_size=1099511627776)
+    cache = {}
+    cnt = 1
+
+    with open(gtFile, 'r', encoding='utf-8') as data:
+        datalist = data.readlines()
+
+    nSamples = len(datalist)
+    for i in range(nSamples):
+        imagePath, label = datalist[i].strip('\n').split('\t')
+        imagePath = os.path.join(inputPath, imagePath)
+
+        # # only use alphanumeric data
+        # if re.search('[^a-zA-Z0-9]', label):
+        #     continue
+
+        if not os.path.exists(imagePath):
+            print('%s does not exist' % imagePath)
+            continue
+        with open(imagePath, 'rb') as f:
+            imageBin = f.read()
+        if checkValid:
+            try:
+                if not checkImageIsValid(imageBin):
+                    print('%s is not a valid image' % imagePath)
+                    continue
+            except:
+                print('error occured', i)
+                with open(outputPath + '/error_image_log.txt', 'a') as log:
+                    log.write('%s-th image data occured error\n' % str(i))
+                continue
+
+        imageKey = 'image-%09d'.encode() % cnt
+        labelKey = 'label-%09d'.encode() % cnt
+        cache[imageKey] = imageBin
+        cache[labelKey] = label.encode()
+
+        if cnt % 1000 == 0:
+            writeCache(env, cache)
+            cache = {}
+            print('Written %d / %d' % (cnt, nSamples))
+        cnt += 1
+    nSamples = cnt-1
+    cache['num-samples'.encode()] = str(nSamples).encode()
+    writeCache(env, cache)
+    print('Created dataset with %d samples' % nSamples)
+
+
+if __name__ == '__main__':
+    fire.Fire(createDataset)

tools/crop_by_word_bb_syn90k.py

@@ -0,0 +1,153 @@
+# Crop by word bounding box
+# Locate script with gt.mat
+# $ python crop_by_word_bb.py
+
+import os
+import re
+import cv2
+import scipy.io as sio
+from itertools import chain
+import numpy as np
+import math
+
+mat_contents = sio.loadmat('gt.mat')
+
+image_names = mat_contents['imnames'][0]
+cropped_indx = 0
+start_img_indx = 0
+gt_file = open('gt_oabc.txt', 'a')
+err_file = open('err_oabc.txt', 'a')
+
+for img_indx in range(start_img_indx, len(image_names)):
+
+
+    # Get image name
+    image_name_new = image_names[img_indx][0]
+    # print(image_name_new)
+    image_name = '/home/yxwang/pytorch/dataset/SynthText/img/'+ image_name_new
+    # print('IMAGE : {}.{}'.format(img_indx, image_name))
+    print('evaluating {} image'.format(img_indx), end='\r')
+    # Get text in image
+    txt = mat_contents['txt'][0][img_indx]
+    txt = [re.split(' \n|\n |\n| ', t.strip()) for t in txt]
+    txt = list(chain(*txt))
+    txt = [t for t in txt if len(t) > 0 ]
+    # print(txt) # ['Lines:', 'I', 'lost', 'Kevin', 'will', 'line', 'and', 'and', 'the', '(and', 'the', 'out', 'you', "don't", 'pkg']
+    # assert 1<0
+
+    # Open image
+    #img = Image.open(image_name)
+    img = cv2.imread(image_name, cv2.IMREAD_COLOR)
+    img_height, img_width, _ = img.shape
+
+    # Validation
+    if len(np.shape(mat_contents['wordBB'][0][img_indx])) == 2:
+        wordBBlen = 1
+    else:
+        wordBBlen = mat_contents['wordBB'][0][img_indx].shape[-1]
+
+    if wordBBlen == len(txt):
+        # Crop image and save
+        for word_indx in range(len(txt)):
+            # print('txt--',txt)
+            txt_temp = txt[word_indx]
+            len_now = len(txt_temp)
+            # txt_temp = re.sub('[^0-9a-zA-Z]+', '', txt_temp)
+            # print('txt_temp-1-',txt_temp)
+            txt_temp = re.sub('[^a-zA-Z]+', '', txt_temp)
+            # print('txt_temp-2-',txt_temp)
+            if len_now - len(txt_temp) != 0:
+                print('txt_temp-2-', txt_temp)
+
+            if len(np.shape(mat_contents['wordBB'][0][img_indx])) == 2:  # only one word (2,4)
+                wordBB = mat_contents['wordBB'][0][img_indx]
+            else:  # many words (2,4,num_words)
+                wordBB = mat_contents['wordBB'][0][img_indx][:, :, word_indx]
+
+            if np.shape(wordBB) != (2, 4):
+                err_log = 'malformed box index: {}\t{}\t{}\n'.format(image_name, txt[word_indx], wordBB)
+                err_file.write(err_log)
+                # print(err_log)
+                continue
+
+            pts1 = np.float32([[wordBB[0][0], wordBB[1][0]],
+                               [wordBB[0][3], wordBB[1][3]],
+                               [wordBB[0][1], wordBB[1][1]],
+                               [wordBB[0][2], wordBB[1][2]]])
+            height = math.sqrt((wordBB[0][0] - wordBB[0][3])**2 + (wordBB[1][0] - wordBB[1][3])**2)
+            width = math.sqrt((wordBB[0][0] - wordBB[0][1])**2 + (wordBB[1][0] - wordBB[1][1])**2)
+
+            # Coord validation check
+            if (height * width) <= 0:
+                err_log = 'empty file : {}\t{}\t{}\n'.format(image_name, txt[word_indx], wordBB)
+                err_file.write(err_log)
+                # print(err_log)
+                continue
+            elif (height * width) > (img_height * img_width):
+                err_log = 'too big box : {}\t{}\t{}\n'.format(image_name, txt[word_indx], wordBB)
+                err_file.write(err_log)
+                # print(err_log)
+                continue
+            else:
+                valid = True
+                for i in range(2):
+                    for j in range(4):
+                        if wordBB[i][j] < 0 or wordBB[i][j] > img.shape[1 - i]:
+                            valid = False
+                            break
+                    if not valid:
+                        break
+                if not valid:
+                    err_log = 'invalid coord : {}\t{}\t{}\t{}\t{}\n'.format(
+                        image_name, txt[word_indx], wordBB, (width, height), (img_width, img_height))
+                    err_file.write(err_log)
+                    # print(err_log)
+                    continue
+
+            pts2 = np.float32([[0, 0],
+                               [0, height],
+                               [width, 0],
+                               [width, height]])
+
+            x_min = np.int(round(min(wordBB[0][0], wordBB[0][1], wordBB[0][2], wordBB[0][3])))
+            x_max = np.int(round(max(wordBB[0][0], wordBB[0][1], wordBB[0][2], wordBB[0][3])))
+            y_min = np.int(round(min(wordBB[1][0], wordBB[1][1], wordBB[1][2], wordBB[1][3])))
+            y_max = np.int(round(max(wordBB[1][0], wordBB[1][1], wordBB[1][2], wordBB[1][3])))
+            # print(x_min, x_max, y_min, y_max)
+            # print(img.shape)
+            # assert 1<0
+            if len(img.shape) == 3:
+                img_cropped = img[ y_min:y_max:1, x_min:x_max:1, :]
+            else:
+                img_cropped = img[ y_min:y_max:1, x_min:x_max:1]
+            dir_name = '/home/yxwang/pytorch/dataset/SynthText/cropped-oabc/{}'.format(image_name_new.split('/')[0])
+            # print('dir_name--',dir_name)
+            if not os.path.exists(dir_name):
+                os.mkdir(dir_name)
+            cropped_file_name = "{}/{}_{}_{}.jpg".format(dir_name, cropped_indx,
+                                                         image_name.split('/')[-1][:-len('.jpg')], word_indx)
+            # print('cropped_file_name--',cropped_file_name)
+            # print('img_cropped--',img_cropped.shape)
+            if img_cropped.shape[0] == 0 or img_cropped.shape[1] == 0:
+                err_log = 'word_box_mismatch : {}\t{}\t{}\n'.format(image_name, mat_contents['txt'][0][
+                    img_indx], mat_contents['wordBB'][0][img_indx])
+                err_file.write(err_log)
+                # print(err_log)
+                continue
+            # print('img_cropped--',img_cropped)
+
+            # img_cropped.save(cropped_file_name)
+            cv2.imwrite(cropped_file_name, img_cropped)
+            cropped_indx += 1
+            gt_file.write('%s\t%s\n' % (cropped_file_name, txt[word_indx]))
+
+            # if cropped_indx>10:
+            #     assert 1<0
+        # assert 1 < 0
+    else:
+        err_log = 'word_box_mismatch : {}\t{}\t{}\n'.format(image_name, mat_contents['txt'][0][
+                                                            img_indx], mat_contents['wordBB'][0][img_indx])
+        err_file.write(err_log)
+        # print(err_log)
+gt_file.close()
+err_file.close()

transforms.py

@@ -0,0 +1,329 @@
+import math
+import numbers
+import random
+
+import cv2
+import numpy as np
+from PIL import Image
+from torchvision import transforms
+from torchvision.transforms import Compose
+
+
+def sample_asym(magnitude, size=None):
+    return np.random.beta(1, 4, size) * magnitude
+
+def sample_sym(magnitude, size=None):
+    return (np.random.beta(4, 4, size=size) - 0.5) * 2 * magnitude
+
+def sample_uniform(low, high, size=None):
+    return np.random.uniform(low, high, size=size)
+
+def get_interpolation(type='random'):
+    if type == 'random':
+        choice = [cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA]
+        interpolation = choice[random.randint(0, len(choice)-1)]
+    elif type == 'nearest': interpolation = cv2.INTER_NEAREST
+    elif type == 'linear': interpolation = cv2.INTER_LINEAR
+    elif type == 'cubic': interpolation = cv2.INTER_CUBIC
+    elif type == 'area': interpolation = cv2.INTER_AREA
+    else: raise TypeError('Interpolation types only nearest, linear, cubic, area are supported!')
+    return interpolation
+
+class CVRandomRotation(object):
+    def __init__(self, degrees=15):
+        assert isinstance(degrees, numbers.Number), "degree should be a single number."
+        assert degrees >= 0, "degree must be positive."
+        self.degrees = degrees
+
+    @staticmethod
+    def get_params(degrees):
+        return sample_sym(degrees)
+    
+    def __call__(self, img):
+        angle = self.get_params(self.degrees)
+        src_h, src_w = img.shape[:2]
+        M = cv2.getRotationMatrix2D(center=(src_w/2, src_h/2), angle=angle, scale=1.0)
+        abs_cos, abs_sin = abs(M[0,0]), abs(M[0,1])
+        dst_w = int(src_h * abs_sin + src_w * abs_cos)
+        dst_h = int(src_h * abs_cos + src_w * abs_sin)
+        M[0, 2] += (dst_w - src_w)/2
+        M[1, 2] += (dst_h - src_h)/2
+        
+        flags = get_interpolation()
+        return cv2.warpAffine(img, M, (dst_w, dst_h), flags=flags, borderMode=cv2.BORDER_REPLICATE)
+
+class CVRandomAffine(object):
+    def __init__(self, degrees, translate=None, scale=None, shear=None):
+        assert isinstance(degrees, numbers.Number), "degree should be a single number."
+        assert degrees >= 0, "degree must be positive."
+        self.degrees = degrees
+
+        if translate is not None:
+            assert isinstance(translate, (tuple, list)) and len(translate) == 2, \
+                "translate should be a list or tuple and it must be of length 2."
+            for t in translate:
+                if not (0.0 <= t <= 1.0):
+                    raise ValueError("translation values should be between 0 and 1")
+        self.translate = translate
+
+        if scale is not None:
+            assert isinstance(scale, (tuple, list)) and len(scale) == 2, \
+                "scale should be a list or tuple and it must be of length 2."
+            for s in scale:
+                if s <= 0:
+                    raise ValueError("scale values should be positive")
+        self.scale = scale
+
+        if shear is not None:
+            if isinstance(shear, numbers.Number):
+                if shear < 0:
+                    raise ValueError("If shear is a single number, it must be positive.")
+                self.shear = [shear]
+            else:
+                assert isinstance(shear, (tuple, list)) and (len(shear) == 2), \
+                    "shear should be a list or tuple and it must be of length 2."
+                self.shear = shear
+        else:
+            self.shear = shear
+            
+    def _get_inverse_affine_matrix(self, center, angle, translate, scale, shear):
+        # https://github.com/pytorch/vision/blob/v0.4.0/torchvision/transforms/functional.py#L717
+        from numpy import sin, cos, tan
+        
+        if isinstance(shear, numbers.Number):
+            shear = [shear, 0]
+
+        if not isinstance(shear, (tuple, list)) and len(shear) == 2:
+            raise ValueError(
+                "Shear should be a single value or a tuple/list containing " +
+                "two values. Got {}".format(shear))
+
+        rot = math.radians(angle)
+        sx, sy = [math.radians(s) for s in shear]
+
+        cx, cy = center
+        tx, ty = translate
+
+        # RSS without scaling
+        a = cos(rot - sy) / cos(sy)
+        b = -cos(rot - sy) * tan(sx) / cos(sy) - sin(rot)
+        c = sin(rot - sy) / cos(sy)
+        d = -sin(rot - sy) * tan(sx) / cos(sy) + cos(rot)
+
+        # Inverted rotation matrix with scale and shear
+        # det([[a, b], [c, d]]) == 1, since det(rotation) = 1 and det(shear) = 1
+        M = [d, -b, 0,
+             -c, a, 0]
+        M = [x / scale for x in M]
+
+        # Apply inverse of translation and of center translation: RSS^-1 * C^-1 * T^-1
+        M[2] += M[0] * (-cx - tx) + M[1] * (-cy - ty)
+        M[5] += M[3] * (-cx - tx) + M[4] * (-cy - ty)
+
+        # Apply center translation: C * RSS^-1 * C^-1 * T^-1
+        M[2] += cx
+        M[5] += cy
+        return M
+
+    @staticmethod
+    def get_params(degrees, translate, scale_ranges, shears, height):        
+        angle = sample_sym(degrees)
+        if translate is not None:
+            max_dx = translate[0] * height
+            max_dy = translate[1] * height
+            translations = (np.round(sample_sym(max_dx)), np.round(sample_sym(max_dy)))
+        else:
+            translations = (0, 0)
+
+        if scale_ranges is not None:
+            scale = sample_uniform(scale_ranges[0], scale_ranges[1])
+        else:
+            scale = 1.0
+
+        if shears is not None:
+            if len(shears) == 1:
+                shear = [sample_sym(shears[0]), 0.]
+            elif len(shears) == 2:
+                shear = [sample_sym(shears[0]), sample_sym(shears[1])]
+        else:
+            shear = 0.0
+
+        return angle, translations, scale, shear
+    
+    
+    def __call__(self, img):
+        src_h, src_w = img.shape[:2]
+        angle, translate, scale, shear = self.get_params(
+                self.degrees, self.translate, self.scale, self.shear, src_h)
+
+        M = self._get_inverse_affine_matrix((src_w/2, src_h/2), angle, (0, 0), scale, shear)
+        M = np.array(M).reshape(2,3)
+        
+        startpoints = [(0, 0), (src_w - 1, 0), (src_w - 1, src_h - 1), (0, src_h - 1)]
+        project = lambda x, y, a, b, c: int(a*x + b*y + c)
+        endpoints = [(project(x, y, *M[0]), project(x, y, *M[1])) for x, y in startpoints]
+        
+        rect = cv2.minAreaRect(np.array(endpoints))
+        bbox = cv2.boxPoints(rect).astype(dtype=np.int)
+        max_x, max_y = bbox[:, 0].max(), bbox[:, 1].max()
+        min_x, min_y = bbox[:, 0].min(), bbox[:, 1].min()
+
+        dst_w = int(max_x - min_x)
+        dst_h = int(max_y - min_y)
+        M[0, 2] += (dst_w - src_w) / 2
+        M[1, 2] += (dst_h - src_h) / 2
+        
+        # add translate
+        dst_w += int(abs(translate[0]))
+        dst_h += int(abs(translate[1]))
+        if translate[0] < 0: M[0, 2] += abs(translate[0])
+        if translate[1] < 0: M[1, 2] += abs(translate[1])
+            
+        flags = get_interpolation()
+        return cv2.warpAffine(img, M, (dst_w , dst_h), flags=flags, borderMode=cv2.BORDER_REPLICATE)
+
+class CVRandomPerspective(object):
+    def __init__(self, distortion=0.5):
+        self.distortion = distortion
+
+    def get_params(self, width, height, distortion):
+        offset_h = sample_asym(distortion * height / 2, size=4).astype(dtype=np.int)
+        offset_w = sample_asym(distortion * width / 2, size=4).astype(dtype=np.int)
+        topleft  = (            offset_w[0],              offset_h[0])
+        topright = (width - 1 - offset_w[1],              offset_h[1])
+        botright = (width - 1 - offset_w[2], height - 1 - offset_h[2])
+        botleft  = (            offset_w[3], height - 1 - offset_h[3])
+
+        startpoints = [(0, 0), (width - 1, 0), (width - 1, height - 1), (0, height - 1)]
+        endpoints = [topleft, topright, botright, botleft]
+        return np.array(startpoints, dtype=np.float32), np.array(endpoints, dtype=np.float32)
+    
+    def __call__(self, img):
+        height, width = img.shape[:2]
+        startpoints, endpoints = self.get_params(width, height, self.distortion)
+        M = cv2.getPerspectiveTransform(startpoints, endpoints)
+
+        # TODO: more robust way to crop image
+        rect = cv2.minAreaRect(endpoints)
+        bbox = cv2.boxPoints(rect).astype(dtype=np.int)
+        max_x, max_y = bbox[:, 0].max(), bbox[:, 1].max()
+        min_x, min_y = bbox[:, 0].min(), bbox[:, 1].min()
+        min_x, min_y = max(min_x, 0), max(min_y, 0)
+
+        flags = get_interpolation() 
+        img = cv2.warpPerspective(img, M, (max_x, max_y), flags=flags, borderMode=cv2.BORDER_REPLICATE)
+        img = img[min_y:, min_x:]
+        return img
+
+class CVRescale(object):
+    
+    def __init__(self, factor=4, base_size=(128, 512)):
+        """ Define image scales using gaussian pyramid and rescale image to target scale.
+        
+        Args:
+            factor: the decayed factor from base size, factor=4 keeps target scale by default.
+            base_size: base size the build the bottom layer of pyramid
+        """
+        if isinstance(factor, numbers.Number):
+            self.factor = round(sample_uniform(0, factor))
+        elif isinstance(factor, (tuple, list)) and len(factor) == 2:
+            self.factor = round(sample_uniform(factor[0], factor[1]))
+        else:
+            raise Exception('factor must be number or list with length 2')
+        # assert factor is valid
+        self.base_h, self.base_w = base_size[:2]
+
+    def __call__(self, img):
+        if self.factor == 0: return img 
+        src_h, src_w = img.shape[:2]
+        cur_w, cur_h = self.base_w, self.base_h        
+        scale_img = cv2.resize(img, (cur_w, cur_h), interpolation=get_interpolation())
+        for _ in range(self.factor): 
+            scale_img = cv2.pyrDown(scale_img)
+        scale_img = cv2.resize(scale_img, (src_w, src_h), interpolation=get_interpolation())
+        return scale_img
+
+class CVGaussianNoise(object):
+    def __init__(self, mean=0, var=20):
+        self.mean = mean
+        if isinstance(var, numbers.Number):
+            self.var = max(int(sample_asym(var)), 1)
+        elif isinstance(var, (tuple, list)) and len(var) == 2:
+            self.var = int(sample_uniform(var[0], var[1]))
+        else:
+            raise Exception('degree must be number or list with length 2')
+        
+    def __call__(self, img):
+        noise = np.random.normal(self.mean, self.var**0.5, img.shape)
+        img = np.clip(img + noise, 0, 255).astype(np.uint8)
+        return img
+
+class CVMotionBlur(object):
+    def __init__(self, degrees=12, angle=90):
+        if isinstance(degrees, numbers.Number):
+            self.degree = max(int(sample_asym(degrees)), 1)
+        elif isinstance(degrees, (tuple, list)) and len(degrees) == 2:
+            self.degree = int(sample_uniform(degrees[0], degrees[1]))
+        else:
+            raise Exception('degree must be number or list with length 2')
+        self.angle = sample_uniform(-angle, angle)
+
+    def __call__(self, img):
+        M = cv2.getRotationMatrix2D((self.degree // 2, self.degree // 2), self.angle, 1)
+        motion_blur_kernel = np.zeros((self.degree, self.degree))
+        motion_blur_kernel[self.degree // 2, :] = 1
+        motion_blur_kernel = cv2.warpAffine(motion_blur_kernel, M, (self.degree, self.degree))
+        motion_blur_kernel = motion_blur_kernel / self.degree
+        img = cv2.filter2D(img, -1, motion_blur_kernel)
+        img = np.clip(img, 0, 255).astype(np.uint8)
+        return img
+
+class CVGeometry(object):
+    def __init__(self, degrees=15, translate=(0.3, 0.3), scale=(0.5, 2.), 
+                       shear=(45, 15), distortion=0.5, p=0.5):
+        self.p = p
+        type_p = random.random()
+        if type_p < 0.33:
+            self.transforms = CVRandomRotation(degrees=degrees)
+        elif type_p < 0.66:
+            self.transforms = CVRandomAffine(degrees=degrees, translate=translate, scale=scale, shear=shear)
+        else:
+            self.transforms = CVRandomPerspective(distortion=distortion)
+
+    def __call__(self, img):
+        if random.random() < self.p:
+            img = np.array(img)
+            return Image.fromarray(self.transforms(img))
+        else: return img
+
+class CVDeterioration(object):
+    def __init__(self, var, degrees, factor, p=0.5):
+        self.p = p
+        transforms = []
+        if var is not None:
+            transforms.append(CVGaussianNoise(var=var))
+        if degrees is not None:
+            transforms.append(CVMotionBlur(degrees=degrees))
+        if factor is not None:
+            transforms.append(CVRescale(factor=factor))
+            
+        random.shuffle(transforms)
+        transforms = Compose(transforms)
+        self.transforms =  transforms
+        
+    def __call__(self, img):
+        if random.random() < self.p:
+            img = np.array(img)
+            return Image.fromarray(self.transforms(img))
+        else: return img
+
+
+class CVColorJitter(object):
+    def __init__(self, brightness=0.5, contrast=0.5, saturation=0.5, hue=0.1, p=0.5):
+        self.p = p
+        self.transforms = transforms.ColorJitter(brightness=brightness, contrast=contrast, 
+                                                 saturation=saturation, hue=hue)
+
+    def __call__(self, img):
+        if random.random() < self.p: return self.transforms(img)
+        else: return img

utils.py

@@ -0,0 +1,304 @@
+import logging
+import os
+import time
+
+import cv2
+import numpy as np
+import torch
+import yaml
+from matplotlib import colors
+from matplotlib import pyplot as plt
+from torch import Tensor, nn
+from torch.utils.data import ConcatDataset
+
+class CharsetMapper(object):
+    """A simple class to map ids into strings.
+
+    It works only when the character set is 1:1 mapping between individual
+    characters and individual ids.
+    """
+
+    def __init__(self,
+                 filename='',
+                 max_length=30,
+                 null_char=u'\u2591'):
+        """Creates a lookup table.
+
+        Args:
+          filename: Path to charset file which maps characters to ids.
+          max_sequence_length: The max length of ids and string.
+          null_char: A unicode character used to replace '<null>' character.
+            the default value is a light shade block '░'.
+        """
+        self.null_char = null_char
+        self.max_length = max_length
+
+        self.label_to_char = self._read_charset(filename)
+        self.char_to_label = dict(map(reversed, self.label_to_char.items()))
+        self.num_classes = len(self.label_to_char)
+ 
+    def _read_charset(self, filename):
+        """Reads a charset definition from a tab separated text file.
+
+        Args:
+          filename: a path to the charset file.
+
+        Returns:
+          a dictionary with keys equal to character codes and values - unicode
+          characters.
+        """
+        import re
+        pattern = re.compile(r'(\d+)\t(.+)')
+        charset = {}
+        self.null_label = 0
+        charset[self.null_label] = self.null_char
+        with open(filename, 'r') as f:
+            for i, line in enumerate(f):
+                m = pattern.match(line)
+                assert m, f'Incorrect charset file. line #{i}: {line}'
+                label = int(m.group(1)) + 1
+                char = m.group(2)
+                charset[label] = char
+        return charset
+
+    def trim(self, text):
+        assert isinstance(text, str)
+        return text.replace(self.null_char, '')
+
+    def get_text(self, labels, length=None, padding=True, trim=False):
+        """ Returns a string corresponding to a sequence of character ids.
+        """
+        length = length if length else self.max_length
+        labels = [l.item() if isinstance(l, Tensor) else int(l) for l in labels]
+        if padding:
+            labels = labels + [self.null_label] * (length-len(labels))
+        text = ''.join([self.label_to_char[label] for label in labels])
+        if trim: text = self.trim(text)
+        return text
+
+    def get_labels(self, text, length=None, padding=True, case_sensitive=False):
+        """ Returns the labels of the corresponding text.
+        """
+        length = length if length else self.max_length
+        if padding:
+            text = text + self.null_char * (length - len(text))
+        if not case_sensitive:
+            text = text.lower()
+        labels = [self.char_to_label[char] for char in text]
+        return labels
+
+    def pad_labels(self, labels, length=None):
+        length = length if length else self.max_length
+
+        return labels + [self.null_label] * (length - len(labels))
+
+    @property
+    def digits(self):
+        return '0123456789'
+
+    @property
+    def digit_labels(self):
+        return self.get_labels(self.digits, padding=False)
+
+    @property
+    def alphabets(self):
+        all_chars = list(self.char_to_label.keys())
+        valid_chars = []
+        for c in all_chars:
+            if c in 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ':
+                valid_chars.append(c)
+        return ''.join(valid_chars)
+
+    @property
+    def alphabet_labels(self):
+        return self.get_labels(self.alphabets, padding=False)
+
+
+class Timer(object):
+    """A simple timer."""
+    def __init__(self):
+        self.data_time = 0.
+        self.data_diff = 0.
+        self.data_total_time = 0.
+        self.data_call = 0
+        self.running_time = 0.
+        self.running_diff = 0.
+        self.running_total_time = 0.
+        self.running_call = 0
+
+    def tic(self):
+        self.start_time = time.time()
+        self.running_time = self.start_time
+
+    def toc_data(self):
+        self.data_time = time.time()
+        self.data_diff = self.data_time - self.running_time
+        self.data_total_time += self.data_diff
+        self.data_call += 1
+
+    def toc_running(self):
+        self.running_time = time.time()
+        self.running_diff = self.running_time - self.data_time
+        self.running_total_time += self.running_diff
+        self.running_call += 1
+
+    def total_time(self):
+        return self.data_total_time + self.running_total_time
+
+    def average_time(self):
+        return  self.average_data_time() + self.average_running_time()
+
+    def average_data_time(self):
+        return self.data_total_time / (self.data_call or 1)
+
+    def average_running_time(self):
+        return self.running_total_time / (self.running_call or 1)
+
+
+class Logger(object):
+    _handle = None
+    _root = None
+
+    @staticmethod
+    def init(output_dir, name, phase):
+        format = '[%(asctime)s %(filename)s:%(lineno)d %(levelname)s {}] ' \
+                '%(message)s'.format(name)
+        logging.basicConfig(level=logging.INFO, format=format)
+
+        try: os.makedirs(output_dir)
+        except: pass
+        config_path = os.path.join(output_dir, f'{phase}.txt')
+        Logger._handle = logging.FileHandler(config_path)
+        Logger._root = logging.getLogger()
+
+    @staticmethod
+    def enable_file():
+        if Logger._handle is None or Logger._root is None:
+            raise Exception('Invoke Logger.init() first!')
+        Logger._root.addHandler(Logger._handle)
+
+    @staticmethod
+    def disable_file():
+        if Logger._handle is None or Logger._root is None:
+            raise Exception('Invoke Logger.init() first!')
+        Logger._root.removeHandler(Logger._handle)
+
+
+class Config(object):
+
+    def __init__(self, config_path, host=True):
+        def __dict2attr(d, prefix=''):
+            for k, v in d.items():
+                if isinstance(v, dict):
+                    __dict2attr(v, f'{prefix}{k}_')
+                else:
+                    if k == 'phase':
+                        assert v in ['train', 'test']
+                    if k == 'stage':
+                        assert v in ['pretrain-vision', 'pretrain-language',
+                                     'train-semi-super', 'train-super']
+                    self.__setattr__(f'{prefix}{k}', v)
+
+        assert os.path.exists(config_path), '%s does not exists!' % config_path
+        with open(config_path) as file:
+            config_dict = yaml.load(file, Loader=yaml.FullLoader)
+        with open('configs/template.yaml') as file:
+            default_config_dict = yaml.load(file, Loader=yaml.FullLoader)
+        __dict2attr(default_config_dict)
+        __dict2attr(config_dict)
+        self.global_workdir = os.path.join(self.global_workdir, self.global_name)
+
+    def __getattr__(self, item):
+        attr = self.__dict__.get(item)
+        if attr is None:
+            attr = dict()
+            prefix = f'{item}_'
+            for k, v in self.__dict__.items():
+                if k.startswith(prefix):
+                    n = k.replace(prefix, '')
+                    attr[n] = v
+            return attr if len(attr) > 0 else None
+        else:
+            return attr
+
+    def __repr__(self):
+        str = 'ModelConfig(\n'
+        for i, (k, v) in enumerate(sorted(vars(self).items())):
+            str += f'\t({i}): {k} = {v}\n'
+        str += ')'
+        return str
+
+def blend_mask(image, mask, alpha=0.5, cmap='jet', color='b', color_alpha=1.0):
+    # normalize mask
+    mask = (mask-mask.min()) / (mask.max() - mask.min() + np.finfo(float).eps)
+    if mask.shape != image.shape:
+        mask = cv2.resize(mask,(image.shape[1], image.shape[0]))
+    # get color map
+    color_map = plt.get_cmap(cmap)
+    mask = color_map(mask)[:,:,:3]
+    # convert float to uint8
+    mask = (mask * 255).astype(dtype=np.uint8)
+
+    # set the basic color
+    basic_color = np.array(colors.to_rgb(color)) * 255 
+    basic_color = np.tile(basic_color, [image.shape[0], image.shape[1], 1]) 
+    basic_color = basic_color.astype(dtype=np.uint8)
+    # blend with basic color
+    blended_img = cv2.addWeighted(image, color_alpha, basic_color, 1-color_alpha, 0)
+    # blend with mask
+    blended_img = cv2.addWeighted(blended_img, alpha, mask, 1-alpha, 0)
+
+    return blended_img
+
+def onehot(label, depth, device=None):
+    """ 
+    Args:
+        label: shape (n1, n2, ..., )
+        depth: a scalar
+
+    Returns:
+        onehot: (n1, n2, ..., depth)
+    """
+    if not isinstance(label, torch.Tensor):
+        label = torch.tensor(label, device=device)
+    onehot = torch.zeros(label.size() + torch.Size([depth]), device=device)
+    onehot = onehot.scatter_(-1, label.unsqueeze(-1), 1)
+
+    return onehot
+
+class MyDataParallel(nn.DataParallel):
+
+    def gather(self, outputs, target_device):
+        r"""
+        Gathers tensors from different GPUs on a specified device
+        (-1 means the CPU).
+        """
+        def gather_map(outputs):
+            out = outputs[0]
+            if isinstance(out, (str, int, float)):
+                return out
+            if isinstance(out, list) and isinstance(out[0], str):
+                return [o for out in outputs for o in out]
+            if isinstance(out, torch.Tensor):
+                return torch.nn.parallel._functions.Gather.apply(target_device, self.dim, *outputs)
+            if out is None:
+                return None
+            if isinstance(out, dict):
+                if not all((len(out) == len(d) for d in outputs)):
+                    raise ValueError('All dicts must have the same number of keys')
+                return type(out)(((k, gather_map([d[k] for d in outputs]))
+                                for k in out))
+            return type(out)(map(gather_map, zip(*outputs)))
+
+        # Recursive function calls like this create reference cycles.
+        # Setting the function to None clears the refcycle.
+        try:
+            res = gather_map(outputs)
+        finally:
+            gather_map = None
+        return res
+
+
+class MyConcatDataset(ConcatDataset):
+    def __getattr__(self, k): 
+        return getattr(self.datasets[0], k)