remove dropout for inference path (#8)

Author: XiaobingSuper

intel_pytorch_extension_py/conf.py

@@ -76,13 +76,14 @@ def get_default_recipe(self, configures):
             #          add                    add
             if len(pre_ops) > 0:
                 for key, value in pre_ops.items():
-                    if value == 'conv2d' or value == 'conv3d':
+                    if value == 'conv2d' or value == 'conv3d' or value == 'linear':
                         default_configures[cur_id]['inputs_quantized'][key] = False
                         break
 
-            # if add pre_op hasn't conv, not need add q, dq for accuracy.
+            # if add pre_op hasn't conv and linear, not need add q, dq for accuracy.
             pre_inputs = pre_ops.values()
-            if cur_op == 'add' and ('conv2d' not in pre_inputs and 'conv3d' not in pre_inputs):
+            if cur_op == 'add' and \
+                    ('conv2d' not in pre_inputs and 'conv3d' not in pre_inputs and 'linear' not in pre_inputs):
                 default_configures[cur_id]['inputs_quantized'][0] = False
                 default_configures[cur_id]['inputs_quantized'][1] = False
 

intel_pytorch_extension_py/fx/__init__.py

@@ -1,2 +1,2 @@
-from .fuse_utils import conv_bn_fuse
+from .fuse_utils import conv_bn_fuse, remove_dropout
 

intel_pytorch_extension_py/fx/fuse_utils.py

@@ -1,8 +1,9 @@
 import torch
+import torch.nn as nn
 import torch.fx as fx
+from torch.fx.node import Argument, Target
 from torch.nn.utils.fusion import fuse_conv_bn_eval
-from typing import Type, Dict, Any, Tuple, Iterable
-import torch
+from typing import Type, Dict, Any, Tuple, Iterable, Optional, List, cast, Callable
 import copy
 
 def _parent_name(target : str) -> Tuple[str, str]:
@@ -63,3 +64,17 @@ def conv_bn_fuse(model: torch.nn.Module, inplace=False) -> torch.nn.Module:
                 new_graph.erase_node(node)
     return fx.GraphModule(fx_model, new_graph)
 
+def remove_dropout(model: nn.Module) -> nn.Module:
+    """
+    Removes all dropout layers from the module.
+    """
+    fx_model = fx.symbolic_trace(model)
+
+    class DropoutRemover(torch.fx.Transformer):
+        def call_module(self, target : Target, args : Tuple[Argument, ...], kwargs : Dict[str, Any]) -> Any:
+            if isinstance(self.submodules[target], nn.Dropout):
+                assert len(args) == 1
+                return args[0]
+            else:
+                return super().call_module(target, args, kwargs)
+    return DropoutRemover(fx_model).transform()