rename module view to lightning view when enable pytorch lightning

tb_plugin/fe/src/app.tsx

@@ -48,7 +48,8 @@ export enum Views {
   Trace = 'Trace',
   Distributed = 'Distributed',
   Memory = 'Memory',
-  Module = 'Module'
+  Module = 'Module',
+  Lightning = 'Lightning'
 }
 
 const ViewNames = {
@@ -58,7 +59,8 @@ const ViewNames = {
   [Views.Trace]: Views.Trace,
   [Views.Distributed]: Views.Distributed,
   [Views.Memory]: Views.Memory,
-  [Views.Module]: Views.Module
+  [Views.Module]: Views.Module,
+  [Views.Lightning]: Views.Lightning
 }
 
 const drawerWidth = 340
@@ -407,6 +409,7 @@ export const App = () => {
         case Views.Memory:
           return <MemoryView run={run} worker={worker} span={span} />
         case Views.Module:
+        case Views.Lightning:
           return <ModuleView run={run} worker={worker} span={span} />
       }
     } else {

tb_plugin/torch_tb_profiler/consts.py

@@ -23,6 +23,7 @@
 DISTRIBUTED_VIEW = View(5, 'distributed', 'Distributed')
 MEMORY_VIEW = View(6, 'memory', 'Memory')
 MODULE_VIEW = View(7, 'module', 'Module')
+LIGHTNING_VIEW = View(8, 'lightning', 'Lightning')
 
 TOOLTIP_GPU_UTIL = \
     'GPU Utilization:\n' \

tb_plugin/torch_tb_profiler/profiler/module_op.py

@@ -4,8 +4,8 @@
 from collections import namedtuple
 from typing import Dict, Generator, Iterable, List, Optional, Set, Tuple, Union
 
-from .node import (DataLoaderNode, ModuleNode, OperatorNode, PLModuleNode,
-                   ProfilerStepNode, is_operator_node)
+from .node import (DataLoaderNode, ModuleNode, OperatorNode, OptimizerNode,
+                   PLModuleNode, ProfilerStepNode, is_operator_node)
 from .trace import BaseEvent, EventTypes, PLModuleEvent, PythonFunctionEvent
 
 
@@ -186,7 +186,8 @@ def _aggregate_modules(modules: Iterable[Union[ModuleNode, PLModuleNode]]) -> Di
 def _get_node_list(tid2tree: Dict[int, OperatorNode], node_class) -> Generator[OperatorNode, None, None]:
     """Get all node with node_class from the operator tree"""
     def traverse_node(node):
-        if type(node) not in (ProfilerStepNode, ModuleNode, OperatorNode, PLModuleNode, DataLoaderNode):
+        # Check OptimizerNode here because in PytorchLightning PLModuleNode is under OptimizerNoder.
+        if type(node) not in (ProfilerStepNode, ModuleNode, OperatorNode, OptimizerNode, PLModuleNode, DataLoaderNode):
             return
 
         if isinstance(node, node_class):

tb_plugin/torch_tb_profiler/profiler/run_generator.py

@@ -69,7 +69,9 @@ def generate_run_profile(self):
 
         profile_run.module_stats = aggegate_module_view(self.profile_data.tid2tree, self.profile_data.events)
         profile_run.pl_module_stats = aggegate_pl_module_view(self.profile_data.tid2tree, self.profile_data.events)
-        if profile_run.module_stats or (profile_run.is_pytorch_lightning and profile_run.pl_module_stats):
+        if profile_run.is_pytorch_lightning and profile_run.pl_module_stats:
+            profile_run.views.append(consts.LIGHTNING_VIEW)
+        elif profile_run.module_stats:
             profile_run.views.append(consts.MODULE_VIEW)
 
         return profile_run