-
- Downloads
Fix non 1:1 mapping between model w. ModuleList and SummaryGraph (#328)
The PyTorch trace mechanism doesn't "see" torch.nn.ModuleList modules (since they don't have a forward function). As a result, the mapping from module names at the Python model definition level to the scope-names at the trace level is not 1:1. This makes it impossible for us to map back from SummaryGraph ops to their respective nn.Modules, which is required for flows like BatchNorm folding and stats fusion in post-training quantization. In #313 we handled this issue specifically in DistillerLSTM, but it makes much more sense to have a generic and automatic solution for this issue, which doesn't require the user to modify the model. This is such a solution. * Implemented DistillerModuleList, a replacement for nn.ModuleList which results in full and unique scope-names * See documentation for this class in summary_graph.py for extensive details on the issue and solution * When generating a SummaryGraph, the model is scanned and all instances of torch.nn.ModuleList are replaced with DistillerModulelist * Add tests for new functionality * Partially revert changes made to DistillerLSTM in commit 43548deb: Keep the refactored _create_cells_list function, but have it create a standard torch.nn.ModuleList (since we're the ModuleList issue automatically now, and no need to confuse users with ad-hoc list implementations
Showing
- distiller/modules/rnn.py 8 additions, 19 deletionsdistiller/modules/rnn.py
- distiller/summary_graph.py 177 additions, 0 deletionsdistiller/summary_graph.py
- examples/word_language_model/manual_lstm_pretrained_stats.yaml 24 additions, 24 deletions...les/word_language_model/manual_lstm_pretrained_stats.yaml
- examples/word_language_model/quantize_lstm.ipynb 4 additions, 4 deletionsexamples/word_language_model/quantize_lstm.ipynb
- tests/test_post_train_quant.py 2 additions, 2 deletionstests/test_post_train_quant.py
- tests/test_summarygraph.py 173 additions, 7 deletionstests/test_summarygraph.py
Loading
Please register or sign in to comment