def find_conv_bn_patterns(self, model, dummy_input):
"""
Find all Conv-BN patterns, used for batch normalization folding
Parameters
----------
model : torch.nn.Module
model to be analyzed.
dummy_input : tupel of torch.tensor
inputs to the model, used for generating the torchscript
"""
if dummy_input is None:
_logger.debug(
"Model inputs are not given, batch normalization folding is disabled"
)
return
graph = build_module_graph(model, dummy_input)
for node_group in graph.nodes_py.nodes_op:
if node_group.op_type in BN_FOLD_OP:
successors = graph.find_successors(node_group.unique_name)
successors = [graph.name_to_node[x] for x in successors]
for successor in successors:
if successor.op_type == 'BatchNorm2d':
self.conv_bn_patterns[node_group.name] = successor.name
def __init__(self, model, dummy_input, masks_file, map_location=None,
batch_dim=0, confidence=8):
assert confidence > 1
# The auto inference will change the values of the parameters in the model
# so we need make a copy before the mask inference
self.ori_state_dict = copy.deepcopy(model.state_dict())
self.bound_model = model
self.inferred_masks = dict() # key: module_name, value: ModuleMasks
self.batch_dim = batch_dim
self.dummy_input, self.device = self._random_model_input(dummy_input, confidence, batch_dim)
self.torch_graph = build_module_graph(model, self.dummy_input)
# dict object to save the auto inferences objects of the submodules
self.auto_inferences = {}
# the index dict to find the corresponding torch._C.Value object
# according to the debug name
# we need the dummy_input to infer the mask automaticlly, so we save
# the indexes from tensor's debugname to the torch._C.Value object.
self.debugname_to_value = {}
# load the mask tensor to the same device with the dummy_input
# self.masks save the mask tensors pruned by the user and the infered
# masks of the others modules
if isinstance(masks_file, (str, Path)) and Path(masks_file).exists():
self.masks = torch.load(
masks_file, map_location if map_location is not None else str(self.device))
elif isinstance(masks_file, dict):
self.masks = masks_file
else:
raise Exception('Please provide the mask or the path of the mask file')
self.constant = {}
# self.internal_result save the internal output of the submodules
self.internal_result = {}
def test_build_module_graph(self):
big_model = BigModel()
g = build_module_graph(big_model, torch.randn(2, 1, 28, 28))
print(g.name_to_node.keys())
leaf_modules = set([
'backbone1.conv1', 'backbone2.bn1', 'backbone2.bn2',
'backbone2.conv1', 'backbone2.conv2', 'backbone2.fc1',
'backbone2.fc2', 'fc3'
])
assert set(g.leaf_modules) == leaf_modules
assert not leaf_modules - set(g.name_to_node.keys())
assert g.find_successors('backbone2.conv1') == ['backbone2.bn1']
assert g.find_successors('backbone2.conv2') == ['backbone2.bn2']
assert g.find_predecessors('backbone2.bn1') == ['backbone2.conv1']
assert g.find_predecessors('backbone2.bn2') == ['backbone2.conv2']
def __init__(self, model, dummy_input, masks_file, map_location=None,
batch_dim=0, confidence=8):
"""
Parameters
----------
model : pytorch model
The model user wants to speed up
dummy_input : pytorch tensor, tuple of tensor, list of tensor
Note: The first dimension of the dummy_input should be the batchsize.
The dummy input for ```jit.trace```, users should put it on the right
device.
masks_file : str
The path of user provided mask file
map_location : str
the device on which masks are placed, same to map_location in ```torch.load```
batch_dim : int
the index of batch dimension in the dummy_input
confidence: the confidence coefficient of the sparsity inference. This value is
actually used as the batchsize of the dummy_input.
"""
assert confidence > 1
# The auto inference will change the values of the parameters in the model
# so we need make a copy before the mask inference
self.ori_state_dict = copy.deepcopy(model.state_dict())
self.bound_model = model
self.inferred_masks = dict() # key: module_name, value: ModuleMasks
self.batch_dim = batch_dim
self.dummy_input, self.device = self._random_model_input(dummy_input, confidence, batch_dim)
self.torch_graph = build_module_graph(model, self.dummy_input)
# dict object to save the auto inferences objects of the submodules
self.auto_inferences = {}
# the index dict to find the corresponding torch._C.Value object
# according to the debug name
# we need the dummy_input to infer the mask automaticlly, so we save
# the indexes from tensor's debugname to the torch._C.Value object.
self.debugname_to_value = {}
# load the mask tensor to the same device with the dummy_input
# self.masks save the mask tensors pruned by the user and the infered
# masks of the others modules
self.masks = torch.load(
masks_file, map_location if map_location is not None else str(self.device))
self.constant = {}
# self.internal_result save the internal output of the submodules
self.internal_result = {}
def __init__(self, model, dummy_input, masks_file, map_location=None):
"""
Parameters
----------
model : pytorch model
The model user wants to speed up
dummy_input : pytorch tensor
The dummy input for ```jit.trace```, users should put it on right device before pass in
masks_file : str
The path of user provided mask file
map_location : str
the device on which masks are placed, same to map_location in ```torch.load```
"""
from nni.common.graph_utils import build_module_graph
self.bound_model = model
self.masks = torch.load(masks_file, map_location)
self.inferred_masks = dict() # key: module_name, value: ModuleMasks
self.dummy_input = dummy_input
self.torch_graph = build_module_graph(model, dummy_input)