Repository: espnet/espnet
Pull request: Fix multi gpu LM issues and add hdf5 LM dataset dump
This PR aims to fix #1075 . I tested this by AN4 and TEDLIUM3. 

I found @kan-bayashi's comment was correct. The memory leak appeared when single GPU performs the frame reduction of loss values from multi GPUs. So I added the `ReduceFramewiseLoss` module to do it in multi GPUs.

In addition, I added HDF5 dump support for LM dataset because some large corpora, e.g., TEDLIUM3, results in 3GB of text dataset and it took 3 mins for preprocessing. After HDF5 dump performed, it only took 30 secs to load that. It really helps me testing this problem, and running rapid try-and-errors in the large LM experiment. (PS. I also tested pickle but it seems to perform very slow and save large bin.)

Problem statement:
multi gpu lm training throws CUDA out of memory exception
There was an out of memory exception thrown when I use lm.yaml and lm_large.yaml in librispeech egs with multi gpu=4.
```
Exception in main training loop: CUDA out of memory. Tried to allocate 19.62 MiB (GPU 2; 22.38 GiB total capacity; 21.01 GiB already allocated; 10.06 MiB free; 634.29 MiB cached)                                   
Traceback (most recent call last):
  File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/chainer/training/trainer.py", line 316, in run                                                                                              
    update()
  File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/chainer/training/updaters/standard_updater.py", line 175, in update                                                                         
    self.update_core()
  File "/mnt/cfs1/asr/users/fanlu/espnet/espnet/lm/pytorch_backend/lm.py", line 277, in update_core
    loss.backward()  # Backprop
  File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/torch/tensor.py", line 102, in backward                                                                                                     
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/torch/autograd/__init__.py", line 90, in backward                                                                                           
    allow_unreachable=True)  # allow_unreachable flag
Will finalize trainer extensions and updater before reraising the exception.
Traceback (most recent call last):
  File "/mnt/cfs1/asr/users/fanlu/espnet/espnet/bin/lm_train.py", line 156, in <module>
    main(sys.argv[1:])
  File "/mnt/cfs1/asr/users/fanlu/espnet/espnet/bin/lm_train.py", line 150, in main
    train(args)
  File "/mnt/cfs1/asr/users/fanlu/espnet/espnet/lm/pytorch_backend/lm.py", line 416, in train
    trainer.run()
  File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/chainer/training/trainer.py", line 349, in run                                                                                              
    six.reraise(*exc_info)
  File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/six.py", line 693, in reraise
    raise value
  File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/chainer/training/trainer.py", line 316, in run                                                                                              
    update()
  File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/chainer/training/updaters/standard_updater.py", line 175, in update                                                                         
    self.update_core()
  File "/mnt/cfs1/asr/users/fanlu/espnet/espnet/lm/pytorch_backend/lm.py", line 277, in update_core
    loss.backward()  # Backprop
  File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/torch/tensor.py", line 102, in backward                                                                                                     
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/torch/autograd/__init__.py", line 90, in backward                                                                                           
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: CUDA out of memory. Tried to allocate 19.62 MiB (GPU 2; 22.38 GiB total capacity; 21.01 GiB already allocated; 10.06 MiB free; 634.29 MiB cached)                                                      
# Accounting: time=746 threads=1
# Finished at Sun Aug 11 14:58:25 CST 2019 with status 1
```

Hints from issue discussion:
ensure your sentence length not too long, may be small than 128/256
There is a same problem is egs tedlium3 too.And there was no problem in single gpu before.So Maybe It's a new bug in multi gpu recipe.
Could you check whether it works with small batchsize?
In multi-gpu case, we automatically increase `batchsize = batchsize * ngpu`.

Yes.I have changed the batchsize from 512 to 256/128/32. There are exists the same problem also.
Thank you. There might be a bug.
We will investigate.
I could reproduce this bug in TED-LIUM3. Because I tested multi GPU LM in AN4, I haven't met this problem. I suspect that this issue is related to chainer trainer because I just simply added `torch.nn.DataParallel` in the standard way. So I will create chainer-less `lm_train.py` for checking that.
Maybe this is related to following part
https://github.com/espnet/espnet/blob/8fdd8e96b0896d97a63ab74ceb1cbc01c7652778/espnet/lm/pytorch_backend/lm.py#L265-L270

The final loss calculation is conducted on only main gpu.
Hmm, if so, this should be a pytorch bug. Maybe we can create the smallest example to reproduce this memory leak in multi GPUs to submit an issue to pytorch.

Commit message: minor update

Patch to review:
diff --git a/egs/tedlium3/asr1/run.sh b/egs/tedlium3/asr1/run.sh
index 9a94eb024ed..26174754b4e 100755
--- a/egs/tedlium3/asr1/run.sh
+++ b/egs/tedlium3/asr1/run.sh
@@ -180,7 +180,8 @@ if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
         --train-label ${lmdatadir}/train.txt \
         --valid-label ${lmdatadir}/valid.txt \
         --resume ${lm_resume} \
-        --dict ${dict}
+        --dict ${dict} \
+        --dump-dataset ${lmdatadir}
 fi
 
 if [ -z ${tag} ]; then
diff --git a/espnet/bin/lm_train.py b/espnet/bin/lm_train.py
index bb820ef1e0e..26f880ae8bc 100755
--- a/espnet/bin/lm_train.py
+++ b/espnet/bin/lm_train.py
@@ -85,6 +85,8 @@ def get_parser():
                         help='dropout probability')
     parser.add_argument('--maxlen', type=int, default=40,
                         help='Batch size is reduced if the input sequence > ML')
+    parser.add_argument('--dump-dataset', type=str, default=None,
+                        help='dump preprocessed dataset as hdf5 if specified')
     return parser
 
 
diff --git a/espnet/lm/lm_utils.py b/espnet/lm/lm_utils.py
index 1c22ebd6f14..7bd83d20992 100644
--- a/espnet/lm/lm_utils.py
+++ b/espnet/lm/lm_utils.py
@@ -6,19 +6,51 @@
 # This code is ported from the following implementation written in Torch.
 # https://github.com/chainer/chainer/blob/master/examples/ptb/train_ptb_custom_loop.py
 
-from __future__ import division
-from __future__ import print_function
-
 import chainer
+import h5py
 import logging
 import numpy as np
 import os
 import random
 import six
+from tqdm import tqdm
 
 from chainer.training import extension
 
 
+def load_dataset(path, label_dict, outdir=None):
+    """Load and save HDF5 dataset for LM
+
+    Args:
+        path (str): The path of an input text dataset file
+        label_dict (dict[str, int]): dictionary that maps token label string to its ID number
+        outdir (str): The path of an output dir
+
+    Returns:
+        list[np.ndarray]: np.int32 IDs of tokens converted by `read_tokens`
+    """
+    if outdir is not None:
+        os.makedirs(outdir, exist_ok=True)
+        filename = outdir + "/" + os.path.basename(path) + ".h5"
+        if os.path.exists(filename):
+            logging.info(f"loading binary dataset: {filename}")
+            f = h5py.File(filename, "r")
+            return f["data"][:], f["n_tokens"][()], f["n_oovs"][()]
+    else:
+        logging.info("skip dump/load HDF5 because the output dir is not specified")
+    logging.info(f"reading text dataset: {path}")
+    ret = read_tokens(path, label_dict)
+    n_tokens, n_oovs = count_tokens(ret, label_dict['<unk>'])
+    if outdir is not None:
+        logging.info(f"saving binary dataset: {filename}")
+        with h5py.File(filename, "w") as f:
+            # http://docs.h5py.org/en/stable/special.html#arbitrary-vlen-data
+            f.create_dataset("data", data=ret, dtype=h5py.special_dtype(vlen=np.int32))
+            f["n_tokens"] = n_tokens
+            f["n_oovs"] = n_oovs
+    return ret, n_tokens, n_oovs
+
+
 def read_tokens(filename, label_dict):
     """Read tokens as a sequence of sentences
 
@@ -29,10 +61,9 @@ def read_tokens(filename, label_dict):
     """
 
     data = []
-    for ln in open(filename, 'rb').readlines():
-        data.append(np.array([label_dict[label]
-                              if label in label_dict else label_dict['<unk>']
-                              for label in ln.decode('utf-8').split()], dtype=np.int32))
+    unk = label_dict['<unk>']
+    for ln in tqdm(open(filename, 'r', encoding='utf-8')):
+        data.append(np.array([label_dict.get(label, unk) for label in ln.split()], dtype=np.int32))
     return data
 
 
diff --git a/espnet/lm/pytorch_backend/lm.py b/espnet/lm/pytorch_backend/lm.py
index 2a9bc87f080..aaac7bafb7d 100644
--- a/espnet/lm/pytorch_backend/lm.py
+++ b/espnet/lm/pytorch_backend/lm.py
@@ -27,6 +27,7 @@
 
 from espnet.lm.lm_utils import compute_perplexity
 from espnet.lm.lm_utils import count_tokens
+from espnet.lm.lm_utils import load_dataset
 from espnet.lm.lm_utils import MakeSymlinkToBestModel
 from espnet.lm.lm_utils import ParallelSentenceIterator
 from espnet.lm.lm_utils import read_tokens
@@ -230,6 +231,52 @@ def concat_examples(batch, device=None, padding=None):
     return x, t
 
 
+class ReduceFramewiseLoss(torch.nn.Module):
+    """Reduce framewise loss values in language modeling
+
+    Args:
+        model (ClassifierWithState): module for computing framewise loss
+
+    Note:
+        PyTorch seems to have memory leak when one GPU compute this after data parallel.
+        If parallel GPUs compute this, it seems to be fine.
+        See also https://github.com/espnet/espnet/issues/1075
+    """
+
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+        self.reporter = model.reporter
+
+    def state_dict(self):
+        return self.model.state_dict()
+
+    def load_state_dict(self, d):
+        self.model.load_state_dict(d)
+
+    def forward(self, x, t):
+        """Reduce framewise loss values
+
+        Args:
+            x (torch.Tensor): input ids. (batch, len)
+            t (torch.Tensor): target ids. (batch, len)
+
+        Returns:
+            torch.Tensor: Reduced loss values and valid element counts
+        """
+        loss = 0
+        count = torch.tensor(0).long()
+        state = None
+        batch_size, sequence_length = x.shape
+        for i in six.moves.range(sequence_length):
+            # Compute the loss at this time step and accumulate it
+            state, loss_batch = self.model(state, x[:, i], t[:, i])
+            non_zeros = torch.sum(x[:, i] != 0, dtype=torch.float)
+            loss += loss_batch * non_zeros
+            count += int(non_zeros)
+        return loss, count.to(loss.device)
+
+
 class BPTTUpdater(training.StandardUpdater):
     """An updater for a pytorch LM
 
@@ -254,24 +301,15 @@ def update_core(self):
         optimizer = self.get_optimizer('main')
         # Progress the dataset iterator for sentences at each iteration.
         batch = train_iter.__next__()
-        x, t = concat_examples(batch, device=self.device, padding=(0, -100))
         # Concatenate the token IDs to matrices and send them to the device
         # self.converter does this job
         # (it is chainer.dataset.concat_examples by default)
-        loss = 0
-        count = 0
-        state = None
-        batch_size, sequence_length = x.shape
-        for i in six.moves.range(sequence_length):
-            # Compute the loss at this time step and accumulate it
-            state, loss_batch = self.model(state, x[:, i], t[:, i])
-            non_zeros = torch.sum(x[:, i] != 0, dtype=torch.float)
-            loss += loss_batch * non_zeros
-            count += int(non_zeros)
-
-        loss = loss.mean()
-        reporter.report({'loss': float(loss.detach())}, optimizer.target)
-        reporter.report({'count': count}, optimizer.target)
+        x, t = concat_examples(batch, device=self.device, padding=(0, -100))
+        loss, count = self.model(x, t)
+        loss = loss.sum()
+        count = count.sum()
+        reporter.report({'loss': float(loss.detach() / count)}, optimizer.target)
+        reporter.report({'count': int(count)}, optimizer.target)
         # update
         self.model.zero_grad()  # Clear the parameter gradients
         loss.backward()  # Backprop
@@ -303,17 +341,12 @@ def evaluate(self):
         with torch.no_grad():
             for batch in copy.copy(val_iter):
                 x, t = concat_examples(batch, device=self.device, padding=(0, -100))
-                state = None
-                for i in six.moves.range(len(x[0])):
-                    state, loss_batch = self.model(state, x[:, i], t[:, i])
-                    non_zeros = torch.sum(x[:, i] != 0, dtype=torch.float)
-                    loss += loss_batch * non_zeros
-                    count += int(non_zeros)
+                loss, count = self.model(x, t)
         self.model.train()
         # report validation loss
         observation = {}
         with reporter.report_scope(observation):
-            reporter.report({'loss': float(loss.mean() / count)}, self.model.reporter)
+            reporter.report({'loss': float(loss.sum() / count.sum())}, self.model.reporter)
         return observation
 
 
@@ -335,11 +368,8 @@ def train(args):
     unk = args.char_list_dict['<unk>']
     eos = args.char_list_dict['<eos>']
     # read tokens as a sequence of sentences
-    train = read_tokens(args.train_label, args.char_list_dict)
-    val = read_tokens(args.valid_label, args.char_list_dict)
-    # count tokens
-    n_train_tokens, n_train_oovs = count_tokens(train, unk)
-    n_val_tokens, n_val_oovs = count_tokens(val, unk)
+    val, n_val_tokens, n_val_oovs = load_dataset(args.valid_label, args.char_list_dict, args.dump_dataset)
+    train, n_train_tokens, n_train_oovs = load_dataset(args.train_label, args.char_list_dict, args.dump_dataset)
     logging.info('#vocab = ' + str(args.n_vocab))
     logging.info('#sentences in the training data = ' + str(len(train)))
     logging.info('#tokens in the training data = ' + str(n_train_tokens))
@@ -361,10 +391,12 @@ def train(args):
     logging.info('#total iterations = ' + str(args.epoch * len(train_iter.batch_indices)))
     # Prepare an RNNLM model
     rnn = RNNLM(args.n_vocab, args.layer, args.unit, args.type, args.dropout_rate)
-    model = ClassifierWithState(rnn)
+    classifier = ClassifierWithState(rnn)
+    reporter = classifier.reporter
+    model = ReduceFramewiseLoss(classifier)
     if args.ngpu > 0:
         model = torch.nn.DataParallel(model, device_ids=list(range(args.ngpu))).cuda()
-        setattr(model, "reporter", model.module.reporter)
+        setattr(model, "reporter", reporter)
         gpu_id = 0
     else:
         gpu_id = -1


Code review comments:
diff_hunk: @@ -85,6 +85,8 @@ def get_parser():
                         help='dropout probability')
     parser.add_argument('--maxlen', type=int, default=40,
                         help='Batch size is reduced if the input sequence > ML')
+    parser.add_argument('--dump-dataset', type=str, default=None,
path: espnet/bin/lm_train.py
text: @user1:
I think it is better to rename or change to bool type to make it clear.

@author:
Maybe `--dump-hdf5-path` is better?

@user1:
Sounds good :)