customizable transform statistics (#2059)

Author: Zenglinxiao

onmt/bin/train.py

@@ -139,7 +139,7 @@ def train(opt):
                 opt, fields, transforms_cls, stride=nb_gpu, offset=device_id)
             producer = mp.Process(target=batch_producer,
                                   args=(train_iter, queues[device_id],
-                                        semaphore, opt,),
+                                        semaphore, opt, device_id),
                                   daemon=True)
             producers.append(producer)
             producers[device_id].start()

onmt/inputters/corpus.py

@@ -123,7 +123,6 @@ def load(self, offset=0, stride=1):
         with exfile_open(self.src, mode='rb') as fs,\
                 exfile_open(self.tgt, mode='rb') as ft,\
                 exfile_open(self.align, mode='rb') as fa:
-            logger.info(f"Loading {repr(self)}...")
             for i, (sline, tline, align) in enumerate(zip(fs, ft, fa)):
                 if (i % stride) == offset:
                     sline = sline.decode('utf-8')
@@ -136,7 +135,7 @@ def load(self, offset=0, stride=1):
                         example['align'] = align.decode('utf-8')
                     yield example
 
-    def __repr__(self):
+    def __str__(self):
         cls_name = type(self).__name__
         return '{}({}, {}, align={})'.format(
             cls_name, self.src, self.tgt, self.align)
@@ -169,19 +168,17 @@ class ParallelCorpusIterator(object):
 
     Args:
         corpus (ParallelCorpus): corpus to iterate;
-        transform (Transform): transforms to be applied to corpus;
-        infinitely (bool): True to iterate endlessly;
+        transform (TransformPipe): transforms to be applied to corpus;
         skip_empty_level (str): security level when encouter empty line;
         stride (int): iterate corpus with this line stride;
         offset (int): iterate corpus with this line offset.
     """
 
-    def __init__(self, corpus, transform, infinitely=False,
+    def __init__(self, corpus, transform,
                  skip_empty_level='warning', stride=1, offset=0):
         self.cid = corpus.id
         self.corpus = corpus
         self.transform = transform
-        self.infinitely = infinitely
         if skip_empty_level not in ['silent', 'warning', 'error']:
             raise ValueError(
                 f"Invalid argument skip_empty_level={skip_empty_level}")
@@ -208,8 +205,11 @@ def _transform(self, stream):
                 yield item
         report_msg = self.transform.stats()
         if report_msg != '':
-            logger.info("Transform statistics for {}:\n{}".format(
-                self.cid, report_msg))
+            logger.info(
+                "* Transform statistics for {}({:.2f}%):\n{}\n".format(
+                    self.cid, 100/self.stride, report_msg
+                )
+            )
 
     def _add_index(self, stream):
         for i, item in enumerate(stream):
@@ -227,24 +227,17 @@ def _add_index(self, stream):
                 continue
             yield item
 
-    def _iter_corpus(self):
+    def __iter__(self):
         corpus_stream = self.corpus.load(
-            stride=self.stride, offset=self.offset)
+            stride=self.stride, offset=self.offset
+        )
         tokenized_corpus = self._tokenize(corpus_stream)
         transformed_corpus = self._transform(tokenized_corpus)
         indexed_corpus = self._add_index(transformed_corpus)
         yield from indexed_corpus
 
-    def __iter__(self):
-        if self.infinitely:
-            while True:
-                _iter = self._iter_corpus()
-                yield from _iter
-        else:
-            yield from self._iter_corpus()
-
 
-def build_corpora_iters(corpora, transforms, corpora_info, is_train=False,
+def build_corpora_iters(corpora, transforms, corpora_info,
                         skip_empty_level='warning', stride=1, offset=0):
     """Return `ParallelCorpusIterator` for all corpora defined in opts."""
     corpora_iters = dict()
@@ -256,7 +249,7 @@ def build_corpora_iters(corpora, transforms, corpora_info, is_train=False,
         transform_pipe = TransformPipe.build_from(corpus_transform)
         logger.info(f"{c_id}'s transforms: {str(transform_pipe)}")
         corpus_iter = ParallelCorpusIterator(
-            corpus, transform_pipe, infinitely=is_train,
+            corpus, transform_pipe,
             skip_empty_level=skip_empty_level, stride=stride, offset=offset)
         corpora_iters[c_id] = corpus_iter
     return corpora_iters
@@ -294,7 +287,7 @@ def build_sub_vocab(corpora, transforms, opts, n_sample, stride, offset):
     sub_counter_src = Counter()
     sub_counter_tgt = Counter()
     datasets_iterables = build_corpora_iters(
-        corpora, transforms, opts.data, is_train=False,
+        corpora, transforms, opts.data,
         skip_empty_level=opts.skip_empty_level,
         stride=stride, offset=offset)
     for c_name, c_iter in datasets_iterables.items():
@@ -380,7 +373,7 @@ def save_transformed_sample(opts, transforms, n_sample=3):
 
     corpora = get_corpora(opts, is_train=True)
     datasets_iterables = build_corpora_iters(
-        corpora, transforms, opts.data, is_train=False,
+        corpora, transforms, opts.data,
         skip_empty_level=opts.skip_empty_level)
     sample_path = os.path.join(
         os.path.dirname(opts.save_data), CorpusName.SAMPLE)

onmt/inputters/dynamic_iterator.py

@@ -6,6 +6,7 @@
 from onmt.inputters.corpus import get_corpora, build_corpora_iters,\
     DatasetAdapter
 from onmt.transforms import make_transforms
+from onmt.utils.logging import logger
 
 
 class MixingStrategy(object):
@@ -47,13 +48,22 @@ class WeightedMixer(MixingStrategy):
 
     def __init__(self, iterables, weights):
         super().__init__(iterables, weights)
-        self._iterators = {
-            ds_name: iter(generator)
-            for ds_name, generator in self.iterables.items()
-        }
+        self._iterators = {}
+        self._counts = {}
+        for ds_name in self.iterables.keys():
+            self._reset_iter(ds_name)
+
+    def _logging(self):
+        """Report corpora loading statistics."""
+        msgs = []
+        for ds_name, ds_count in self._counts.items():
+            msgs.append(f"\t\t\t* {ds_name}: {ds_count}")
+        logger.info("Weighted corpora loaded so far:\n"+"\n".join(msgs))
 
     def _reset_iter(self, ds_name):
         self._iterators[ds_name] = iter(self.iterables[ds_name])
+        self._counts[ds_name] = self._counts.get(ds_name, 0) + 1
+        self._logging()
 
     def _iter_datasets(self):
         for ds_name, ds_weight in self.weights.items():
@@ -144,8 +154,7 @@ def from_opts(cls, corpora, transforms, fields, opts, is_train,
 
     def _init_datasets(self):
         datasets_iterables = build_corpora_iters(
-            self.corpora, self.transforms,
-            self.corpora_info, self.is_train,
+            self.corpora, self.transforms, self.corpora_info,
             skip_empty_level=self.skip_empty_level,
             stride=self.stride, offset=self.offset)
         self.dataset_adapter = DatasetAdapter(self.fields, self.is_train)

onmt/opts.py

@@ -26,6 +26,10 @@ def _add_logging_opts(parser, is_train=True):
               action=StoreLoggingLevelAction,
               choices=StoreLoggingLevelAction.CHOICES,
               default="0")
+    group.add('--verbose', '-verbose', action="store_true",
+              help='Print data loading and statistics for all process'
+              '(default only log the first process shard)' if is_train
+              else 'Print scores and predictions for each sentence')
 
     if is_train:
         group.add('--report_every', '-report_every', type=int, default=50,
@@ -44,8 +48,6 @@ def _add_logging_opts(parser, is_train=True):
                        "This is also the name of the run.")
     else:
         # Options only during inference
-        group.add('--verbose', '-verbose', action="store_true",
-                  help='Print scores and predictions for each sentence')
         group.add('--attn_debug', '-attn_debug', action="store_true",
                   help='Print best attn for each word')
         group.add('--align_debug', '-align_debug', action="store_true",
@@ -75,7 +77,7 @@ def _add_dynamic_corpus_opts(parser, build_vocab_only=False):
               help="Security level when encounter empty examples."
                    "silent: silently ignore/skip empty example;"
                    "warning: warning when ignore/skip empty example;"
-                   "error: raise error & stop excution when encouter empty.)")
+                   "error: raise error & stop execution when encouter empty.")
     group.add("-transforms", "--transforms", default=[], nargs="+",
               choices=AVAILABLE_TRANSFORMS.keys(),
               help="Default transform pipeline to apply to data. "

onmt/tests/test_transform.py

@@ -5,7 +5,12 @@
 import yaml
 import math
 from argparse import Namespace
-from onmt.transforms import get_transforms_cls, get_specials, make_transforms
+from onmt.transforms import (
+    get_transforms_cls,
+    get_specials,
+    make_transforms,
+    TransformPipe,
+)
 from onmt.transforms.bart import BARTNoising
 
 
@@ -51,6 +56,47 @@ def test_transform_specials(self):
         self.assertEqual(specials, specials_expected)
 
 
+    def test_transform_pipe(self):
+        # 1. Init first transform in the pipe
+        prefix_cls = get_transforms_cls(["prefix"])["prefix"]
+        corpora = yaml.safe_load("""
+            trainset:
+                path_src: data/src-train.txt
+                path_tgt: data/tgt-train.txt
+                transforms: [prefix, filtertoolong]
+                weight: 1
+                src_prefix: "｟_pf_src｠"
+                tgt_prefix: "｟_pf_tgt｠"
+        """)
+        opt = Namespace(data=corpora, seed=-1)
+        prefix_transform = prefix_cls(opt)
+        prefix_transform.warm_up()
+        # 2. Init second transform in the pipe
+        filter_cls = get_transforms_cls(["filtertoolong"])["filtertoolong"]
+        opt = Namespace(src_seq_length=4, tgt_seq_length=4)
+        filter_transform = filter_cls(opt)
+        # 3. Sequential combine them into a transform pipe
+        transform_pipe = TransformPipe.build_from(
+            [prefix_transform, filter_transform]
+        )
+        ex = {
+            "src": ["Hello", ",", "world", "."],
+            "tgt": ["Bonjour", "le", "monde", "."],
+        }
+        # 4. apply transform pipe for example
+        ex_after = transform_pipe.apply(
+            copy.deepcopy(ex), corpus_name="trainset"
+        )
+        # 5. example after the pipe exceed the length limit, thus filtered
+        self.assertIsNone(ex_after)
+        # 6. Transform statistics registed (here for filtertoolong)
+        self.assertTrue(len(transform_pipe.statistics.observables) > 0)
+        msg = transform_pipe.statistics.report()
+        self.assertIsNotNone(msg)
+        # 7. after report, statistics become empty as a fresh start
+        self.assertTrue(len(transform_pipe.statistics.observables) == 0)
+
+
 class TestMiscTransform(unittest.TestCase):
     def test_prefix(self):
         prefix_cls = get_transforms_cls(["prefix"])["prefix"]

onmt/transforms/misc.py

@@ -1,6 +1,17 @@
 from onmt.utils.logging import logger
 from onmt.transforms import register_transform
-from .transform import Transform
+from .transform import Transform, ObservableStats
+
+
+class FilterTooLongStats(ObservableStats):
+    """Runing statistics for FilterTooLongTransform."""
+    __slots__ = ["filtered"]
+
+    def __init__(self):
+        self.filtered = 1
+
+    def update(self, other: "FilterTooLongStats"):
+        self.filtered += other.filtered
 
 
 @register_transform(name='filtertoolong')
@@ -28,7 +39,7 @@ def apply(self, example, is_train=False, stats=None, **kwargs):
         if (len(example['src']) > self.src_seq_length or
                 len(example['tgt']) > self.tgt_seq_length):
             if stats is not None:
-                stats.filter_too_long()
+                stats.update(FilterTooLongStats())
             return None
         else:
             return example

onmt/transforms/sampling.py

@@ -3,7 +3,7 @@
 import numpy as np
 from onmt.constants import DefaultTokens
 from onmt.transforms import register_transform
-from .transform import Transform
+from .transform import Transform, ObservableStats
 
 
 class HammingDistanceSampling(object):
@@ -44,6 +44,20 @@ def _set_seed(self, seed):
         random.seed(seed)
 
 
+class SwitchOutStats(ObservableStats):
+    """Runing statistics for counting tokens being switched out."""
+
+    __slots__ = ["changed", "total"]
+
+    def __init__(self, changed: int, total: int):
+        self.changed = changed
+        self.total = total
+
+    def update(self, other: "SwitchOutStats"):
+        self.changed += other.changed
+        self.total += other.total
+
+
 @register_transform(name='switchout')
 class SwitchOutTransform(HammingDistanceSamplingTransform):
     """
@@ -81,7 +95,7 @@ def _switchout(self, tokens, vocab, stats=None):
         for i in chosen_indices:
             tokens[i] = self._sample_replace(vocab, reject=tokens[i])
         if stats is not None:
-            stats.switchout(n_switchout=n_chosen, n_total=len(tokens))
+            stats.update(SwitchOutStats(n_chosen, len(tokens)))
         return tokens
 
     def apply(self, example, is_train=False, stats=None, **kwargs):
@@ -98,6 +112,20 @@ def _repr_args(self):
         return '{}={}'.format('switchout_temperature', self.temperature)
 
 
+class TokenDropStats(ObservableStats):
+    """Runing statistics for counting tokens being switched out."""
+
+    __slots__ = ["dropped", "total"]
+
+    def __init__(self, dropped: int, total: int):
+        self.dropped = dropped
+        self.total = total
+
+    def update(self, other: "TokenDropStats"):
+        self.dropped += other.dropped
+        self.total += other.total
+
+
 @register_transform(name='tokendrop')
 class TokenDropTransform(HammingDistanceSamplingTransform):
     """Random drop tokens from sentence."""
@@ -126,7 +154,7 @@ def _token_drop(self, tokens, stats=None):
         out = [tok for (i, tok) in enumerate(tokens)
                if i not in chosen_indices]
         if stats is not None:
-            stats.token_drop(n_dropped=n_chosen, n_total=n_items)
+            stats.update(TokenDropStats(n_chosen, n_items))
         return out
 
     def apply(self, example, is_train=False, stats=None, **kwargs):
@@ -141,6 +169,20 @@ def _repr_args(self):
         return '{}={}'.format('tokendrop_temperature', self.temperature)
 
 
+class TokenMaskStats(ObservableStats):
+    """Runing statistics for counting tokens being switched out."""
+
+    __slots__ = ["masked", "total"]
+
+    def __init__(self, masked: int, total: int):
+        self.masked = masked
+        self.total = total
+
+    def update(self, other: "TokenMaskStats"):
+        self.masked += other.masked
+        self.total += other.total
+
+
 @register_transform(name='tokenmask')
 class TokenMaskTransform(HammingDistanceSamplingTransform):
     """Random mask tokens from src sentence."""
@@ -175,7 +217,7 @@ def _token_mask(self, tokens, stats=None):
         for i in chosen_indices:
             tokens[i] = self.MASK_TOK
         if stats is not None:
-            stats.token_mask(n_masked=n_chosen, n_total=len(tokens))
+            stats.update(TokenDropStats(n_chosen, len(tokens)))
         return tokens
 
     def apply(self, example, is_train=False, stats=None, **kwargs):