Repository: espnet/espnet
Pull request: fix #1751
register ctc loss_fn as float function
@ShigekiKarita

Problem statement:
built in ctc conflict with apex
```
# asr_train.py --config conf/tuning/train_pytorch_transformer_medium.yaml --train-dtype O1 --ctc_type builtin --ngpu 4 --backend pytorch --outdir exp/train_pytorch_129.d_resume/results --tensorboard-dir tensorboard/train_pytorch_129.d_resume --debugmode 1 --dict data/lang_char/train_unigram5000_units.txt --debugdir exp/train_pytorch_129.d_resume --minibatches 0 --verbose 0 --resume exp/train_pytorch_129.d/results/snapshot.ep.2 --train-json dump/train/deltafalse/data_unigram5000.json --valid-json dump/dev/deltafalse/data_unigram5000.json 
# Started at Tue Mar 31 12:42:36 CST 2020
#
2020-03-31 12:42:37,647 (asr_train:299) WARNING: Skip DEBUG/INFO messages
Selected optimization level O1:  Insert automatic casts around Pytorch functions and Tensor methods.

Defaults for this optimization level are:
enabled                : True
opt_level              : O1
cast_model_type        : None
patch_torch_functions  : True
keep_batchnorm_fp32    : None
master_weights         : None
loss_scale             : dynamic
Processing user overrides (additional kwargs that are not None)...
After processing overrides, optimization options are:
enabled                : True
opt_level              : O1
cast_model_type        : None
patch_torch_functions  : True
keep_batchnorm_fp32    : None
master_weights         : None
loss_scale             : dynamic
[JException in main training loop: "ctc_loss" not implemented for 'torch.cuda.HalfTensor'
Traceback (most recent call last):
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/chainer/training/trainer.py", line 316, in run
    update()
  File "/data4/tanghaoyu/espnet/espnet/asr/pytorch_backend/asr.py", line 223, in update
    self.update_core()
  File "/data4/tanghaoyu/espnet/espnet/asr/pytorch_backend/asr.py", line 192, in update_core
    loss = data_parallel(self.model, x, range(self.ngpu)).mean() / self.accum_grad
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py", line 188, in data_parallel
    outputs = parallel_apply(replicas, inputs, module_kwargs, used_device_ids)
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py", line 83, in parallel_apply
    raise output
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py", line 59, in _worker
    output = module(*input, **kwargs)
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/torch/nn/modules/module.py", line 489, in __call__
    result = self.forward(*input, **kwargs)
  File "/data4/tanghaoyu/espnet/espnet/nets/pytorch_backend/e2e_asr_transformer.py", line 190, in forward
    loss_ctc = self.ctc(hs_pad.view(batch_size, -1, self.adim), hs_len, ys_pad)
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/torch/nn/modules/module.py", line 489, in __call__
    result = self.forward(*input, **kwargs)
  File "/data4/tanghaoyu/espnet/espnet/nets/pytorch_backend/ctc.py", line 96, in forward
    self.loss = to_device(self, self.loss_fn(ys_hat, ys_true, hlens, olens)).to(dtype=dtype)
  File "/data4/tanghaoyu/espnet/espnet/nets/pytorch_backend/ctc.py", line 50, in loss_fn
    loss = self.ctc_loss(th_pred, th_target, th_ilen, th_olen)
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/torch/nn/modules/module.py", line 489, in __call__
    result = self.forward(*input, **kwargs)
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/torch/nn/modules/loss.py", line 1248, in forward
    return F.ctc_loss(log_probs, targets, input_lengths, target_lengths, self.blank, self.reduction)
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/apex/amp/wrap.py", line 53, in wrapper
    return orig_fn(*args, **kwargs)
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/torch/nn/functional.py", line 1732, in ctc_loss
    return torch.ctc_loss(log_probs, targets, input_lengths, target_lengths, blank, _Reduction.get_enum(reduction))
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/apex/amp/wrap.py", line 53, in wrapper
    return orig_fn(*args, **kwargs)
Will finalize trainer extensions and updater before reraising the exception.
Traceback (most recent call last):
  File "/data4/tanghaoyu/espnet/egs/BIGOenglish/asr1/../../../espnet/bin/asr_train.py", line 368, in <module>
    main(sys.argv[1:])
  File "/data4/tanghaoyu/espnet/egs/BIGOenglish/asr1/../../../espnet/bin/asr_train.py", line 355, in main
    train(args)
  File "/data4/tanghaoyu/espnet/espnet/asr/pytorch_backend/asr.py", line 636, in train
    trainer.run()
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/chainer/training/trainer.py", line 349, in run
    six.reraise(*exc_info)
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/six.py", line 696, in reraise
    raise value
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/chainer/training/trainer.py", line 316, in run
    update()
  File "/data4/tanghaoyu/espnet/espnet/asr/pytorch_backend/asr.py", line 223, in update
    self.update_core()
  File "/data4/tanghaoyu/espnet/espnet/asr/pytorch_backend/asr.py", line 192, in update_core
    loss = data_parallel(self.model, x, range(self.ngpu)).mean() / self.accum_grad
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py", line 188, in data_parallel
    outputs = parallel_apply(replicas, inputs, module_kwargs, used_device_ids)
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py", line 83, in parallel_apply
    raise output
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py", line 59, in _worker
    output = module(*input, **kwargs)
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/torch/nn/modules/module.py", line 489, in __call__
    result = self.forward(*input, **kwargs)
  File "/data4/tanghaoyu/espnet/espnet/nets/pytorch_backend/e2e_asr_transformer.py", line 190, in forward
    loss_ctc = self.ctc(hs_pad.view(batch_size, -1, self.adim), hs_len, ys_pad)
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/torch/nn/modules/module.py", line 489, in __call__
    result = self.forward(*input, **kwargs)
  File "/data4/tanghaoyu/espnet/espnet/nets/pytorch_backend/ctc.py", line 96, in forward
    self.loss = to_device(self, self.loss_fn(ys_hat, ys_true, hlens, olens)).to(dtype=dtype)
  File "/data4/tanghaoyu/espnet/espnet/nets/pytorch_backend/ctc.py", line 50, in loss_fn
    loss = self.ctc_loss(th_pred, th_target, th_ilen, th_olen)
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/torch/nn/modules/module.py", line 489, in __call__
    result = self.forward(*input, **kwargs)
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/torch/nn/modules/loss.py", line 1248, in forward
    return F.ctc_loss(log_probs, targets, input_lengths, target_lengths, self.blank, self.reduction)
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/apex/amp/wrap.py", line 53, in wrapper
    return orig_fn(*args, **kwargs)
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/torch/nn/functional.py", line 1732, in ctc_loss
    return torch.ctc_loss(log_probs, targets, input_lengths, target_lengths, blank, _Reduction.get_enum(reduction))
  File "/data4/tanghaoyu/espnet/tools/venv/lib/python3.7/site-packages/apex/amp/wrap.py", line 53, in wrapper
    return orig_fn(*args, **kwargs)
RuntimeError: "ctc_loss" not implemented for 'torch.cuda.HalfTensor'
# Accounting: time=112 threads=1
# Ended (code 1) at Tue Mar 31 12:44:28 CST 2020, elapsed time 112 seconds
```

Hints from issue discussion:
Which version of pytorch are you using?
I'm expecting CTC loss must be normal precision to avoid over/underflow and it would be intended normal behavior.
Also, was it OK for warpctc?
1.0.1 pytorch
It is ok with warpctc.
So I think one of ctc function should be registed as float function in amp module, How do you think fix by this idea?
@sw005320  
I see, but I'm bit confused because the loss (including CTC function) is used in an `amp` module
https://github.com/espnet/espnet/blob/95ecf5f15fe2577ec649363e9700cd5721885c02/espnet/asr/pytorch_backend/asr.py#L193-L198
the line from 193, 194, 197, 198 seems the standard code from apex example.
Cast the ctc inputs to float32 explicitly. I did the trick for warp ctc here because it also did not support float16/64. Maybe you can update this if statement to `if self.ctc_type == "warpctc" or dtype == torch.float16:`.
https://github.com/espnet/espnet/blob/a97a270aa581d7788bf4d71d108a1e5bd3d23f7d/espnet/nets/pytorch_backend/ctc.py#L90-L92
Also we should update this test to include the built-in ctc mode in float16 https://github.com/espnet/espnet/blob/master/test/test_train_dtype.py
I register the loss_fn as a float function in amp module, But your advice seems more efficient, I will use that implement.
In addition to the efficiency, I prefer the cast because pytorch recently merged `amp` module in master, and its API has changed from NVIDIA APEX. When you find it works, can you send us PR? That would be super helpful for us.
Yes, it works, I can make a PR for it. but you might find a more smart way to implement it as well.
Thanks!

2020年4月6日(月) 18:15 Charlie_Tang <notifications@github.com>:

> Yes, it works, I can make a PR for it. but you might find a more smart way
> to implement it as well.
>
> —
> You are receiving this because you commented.
> Reply to this email directly, view it on GitHub
> <https://github.com/espnet/espnet/issues/1751#issuecomment-609674382>, or
> unsubscribe
> <https://github.com/notifications/unsubscribe-auth/ABTOZ3WCKC4G6WFIKUMDMZLRLGMS5ANCNFSM4LXHHNRA>
> .
>

Commit message: fix #1751

register ctc loss_fn as float function

Patch to review:
diff --git a/espnet/asr/pytorch_backend/asr.py b/espnet/asr/pytorch_backend/asr.py
index 0fd5e8e5447..c4b2e787ab0 100644
--- a/espnet/asr/pytorch_backend/asr.py
+++ b/espnet/asr/pytorch_backend/asr.py
@@ -447,6 +447,12 @@ def train(args):
         else:
             model, optimizer = amp.initialize(model, optimizer, opt_level=args.train_dtype)
         use_apex = True
+
+        if args.ctc_type == "builtin":
+            from espnet.nets.pytorch_backend.ctc import CTC
+            amp.register_float_function(CTC, "loss_fn")
+            amp.init()
+            logging.warning('register pytorch builtin ctc as float function')
     else:
         use_apex = False
 


Code review comments:
diff_hunk: @@ -447,6 +447,12 @@ def train(args):
         else:
             model, optimizer = amp.initialize(model, optimizer, opt_level=args.train_dtype)
         use_apex = True
+
+        if args.ctc_type == "builtin":
+            from espnet.nets.pytorch_backend.ctc import CTC
+            amp.register_float_function(CTC, "loss_fn")
+            amp.init()
+            logging.warning('register pytorch builtin ctc as float function')
path: espnet/asr/pytorch_backend/asr.py
text: @user1:
```suggestion
        from espnet.nets.pytorch_backend.ctc import CTC
        amp.register_float_function(CTC, "loss_fn")
        amp.init()
        logging.warning('register ctc as float function')
```
Maybe we do not need to limit this to "builtin" because warpctc does not support float16 too.

@author:
then the loss_fn of CTC could be simplified, isn't it?

@user1:
Yes if the apex.amp is always enabled. However, I would like to keep the ctc module because of  the disabled case and the latest pytorch amp.