Skip to content
Snippets Groups Projects
user avatar
Guy Jacob authored
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
b614330c
History
Name Last commit Last update
..