NeMo: a toolkit for building AI applications using Neural Modules
Paper · 2019 · Oleksii Kuchaiev, Jason Li, Huyen Nguyen et al.
Summary
Deep learning code is usually built around raw tensors whose axes have no declared meaning. Architecture, training logic and infrastructure also tend to end up tangled in one script, so combining components from different authors is fragile and requires debugging. The authors present NeMo, an open-source Python toolkit that applies standard software-engineering practice to this problem. Its basic unit is the neural module, a reusable, parameterized block such as an encoder, decoder, data layer or loss. Each module's input and output ports carry neural types that describe what every axis means and how large it is. Users connect modules into a graph. NeMo checks type compatibility while the graph is built and computes nothing until a train, eval or infer action is called. The backend is currently PyTorch, with multi-GPU training and mixed precision. Pre-built collections cover speech recognition and NLP. The examples show swapping a CTC decoder for an attention decoder around the same Jasper encoder, and training a Transformer-BIG translation model. The approach makes model parts interchangeable and exposes silent shape bugs early.
Key points
- NeMo has two layers. NeMo Core provides the module base class and type system. The collections, nemo_asr and nemo_nlp, supply ready-made modules.
- A Neural Type maps each tensor axis to an AxisType built from a semantic tag and an optional size. Comparing two types returns one of SAME, LESS, GREATER, DIM_INCOMPATIBLE, TRANSPOSE_SAME or INCOMPATIBLE.
- nemo_asr supports two model families:
- CTC models, such as a Jasper-like model with JasperEncoder, JasperDecoderForCTC and CTCLossNM.
- Attention-based seq2seq models similar to LAS, using JasperRNNConnector, DecoderRNN and SequenceLoss.
- nemo_nlp covers NMT, language modeling, sentence classification, ASR correction, and intent classification performed jointly with slot filling. Via pytorch-transformers, the collection additionally handles pre-training and fine-tuning of BERT.
- A Transformer-BIG model trained on WMT16 English-German for about 15 hours on one 8-GPU machine reached 29.2 BLEU / 28.5 SacreBLEU on newstest2014.
Interesting ideas
- I1. Semantic axis types with inheritance. Each axis gets a meaning tag, and tags form an is-a hierarchy, so a port expecting a general Embedding can accept a WordEmbedding but not the reverse. Any data pipeline that passes arrays between components could use this to reject incompatible inputs that merely happen to have the right shape.
- I2. Detecting transposed axes as a distinct error. If one component emits time-first tensors and another expects batch-first, frameworks usually keep running and the model just fails to converge. Reporting a "transposed but otherwise identical" result at connection time turns this hidden bug into an immediate, explainable error.
- I3. Lazy graph assembly with checks before execution. Building the full module graph first and validating every connection before any computation means configuration mistakes surface in seconds rather than after hours of training. This is useful wherever runs are expensive.
- I4. Configuration only at construction time, with no side effects. Modules depend only on parameters given when they are created and interact only through explicit ports. The network structure then becomes data rather than code, which simplifies experiment tracking, hyperparameter search and reuse of components across projects.
- I5. Cheap wrapping and adapter modules. An existing PyTorch nn.Module becomes a neural module once its ports are declared. A small connector module can bridge a dimensionality gap, as the authors do to reuse the Jasper encoder with an RNN decoder. This lowers the cost of adopting the system and of mixing pre-trained parts from different sources.
Fields, topics and concepts
📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksArts and Humanities · 17,450 worksLanguage and Linguistics · 17,296 works
Cite
Cite as Oleksii Kuchaiev, Jason Li, Huyen Nguyen, Oleksii Hrinchuk, Ryan Leary, Boris Ginsburg, Samuel Kriman, Stanislav Beliaev, Vitaly Lavrukhin, Jack Cook, Patrice Castonguay, Mariya Popova, Jocelyn Huang and Jonathan M. Cohen. 2019. NeMo: a toolkit for building AI applications using Neural Modules. arXiv:1909.09577.
@misc{kuchaiev2019nemo,
title = {{NeMo: a toolkit for building AI applications using Neural Modules}},
author = {Oleksii Kuchaiev, Jason Li, Huyen Nguyen, Oleksii Hrinchuk, Ryan Leary, Boris Ginsburg, Samuel Kriman, Stanislav Beliaev, Vitaly Lavrukhin, Jack Cook, Patrice Castonguay, Mariya Popova, Jocelyn Huang and Jonathan M. Cohen},
year = {2019},
doi = {10.48550/arxiv.1909.09577},
eprint = {1909.09577},
archivePrefix = {arXiv},
url = {https://doi.org/10.48550/arxiv.1909.09577}
}References
- Pytorchpipe. https://github.com/ibm/pytorchpipe. Accessed: 2019-08-19. ↗
- A pytorch extension: Tools for easy mixed precision and distributed training in pytorch. https://github.com/NVIDIA/apex. Accessed: 2019-08-19. ↗
- Pytorch ignite. https://pytorch.org/ignite/. Accessed: 2019-08-19. ↗
- Pytorch lightning. https://github.com/williamFalcon/pytorch-lightning. cessed: 2019-08-19. ↗
- Ludwig. http://ludwig.ai. Accessed: 2019-08-19. ↗
- Tensor considered harmful. http://nlp.seas.harvard.edu/NamedTensor. Accessed: 2019-08-19. ↗
- pytorch-transformers from huggingface. https://github.com/huggingface/pytorch-transformers. Accessed: 2019-08-19. ↗
- Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenber… find
- Denny Britz, Anna Goldie, Thang Luong, and Quoc Le. Massive Exploration of Neural Machine Translation Architectures. ArXiv e-prints, March 2017.
- William Chan, Navdeep Jaitly, Quoc Le, and Oriol Vinyals. Listen, attend and spell: A neural network for large vocabulary conversational speech recognition. In 2016 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 4960–4964. IEEE, 2016. find
- François Chollet et al. Keras. https://keras.io, 2015. ↗
- Guillaume Klein, Yoon Kim, Yuntian Deng, Jean Senellart, and Alexander Rush. OpenNMT: Open-source toolkit for neural machine translation. In Proceedings of ACL 2017, System Demonstrations, pages 67–72, Vancouver, Canada, July 2017. Association for Computational Linguistics. URL https://www.aclweb.or… ↗
- Oleksii Kuchaiev, Boris Ginsburg, Igor Gitman, Vitaly Lavrukhin, Carl Case, and Paulius Micikevicius. Openseq2seq: extensible toolkit for distributed and mixed precision training of sequence-to-sequence models. arXiv e-prints arXiv:1805.10387, 2018. find
- J. Li, V. Lavrukhin, B. Ginsburg, R. Leary, O. Kuchaiev, J. M. Cohen, H. Nguyen, and R. T. Gaddei. Jasper: An end-to-end convolutional neural acoustic model. arXiv e-prints arXiv:1904.03288, 2019. find
- Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory F. Diamos, Erich Elsen, David García, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. Mixed precision training. ICLR, 2017. find
- Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of NAACL-HLT 2019: Demonstrations, 2019.
- Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017. Ac- find
- Malcolm Reynolds, Gabriel Barth-Maron, Frederic Besse, Diego de Las Casas, Andreas Fidjeland, Tim Green, Adrià Puigdomènech, Sébastien Racanière, Jack Rae, and Fabio Viola. Open sourcing Sonnet - a new library for constructing neural networks. https://deepmind.com/blog/open-sourcing-sonnet/, 2017. ↗
- Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. arXiv preprint arXiv: 1706.03762, 2017.
- Ashish Vaswani, Samy Bengio, Eugene Brevdo, Francois Chollet, Aidan N. Gomez, Stephan Gouws, Llion Jones, Łukasz Kaiser, Nal Kalchbrenner, Niki Parmar, Ryan Sepassi, Noam Shazeer, and Jakob Uszkoreit. Tensor2tensor for neural machine translation. CoRR, abs/1803.07416, 2018. URL http://arxiv.org/abs/… ↗
Updated
This page in Markdown → https://md.research.elamur.ai/free/en/papers/kuchaiev-2019-nemo-toolkit-building-ai-applications/