Skip to content
Snippets Groups Projects
  • Guy Jacob's avatar
    b614330c
    Fix non 1:1 mapping between model w. ModuleList and SummaryGraph (#328) · b614330c
    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
    Fix non 1:1 mapping between model w. ModuleList and SummaryGraph (#328)
    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
test NaN GiB