def process_inline_mutation(model: Model) -> Optional[List[Mutator]]:
applied_mutators = []
ic_nodes = _group_by_label(model.get_nodes_by_type('__torch__.nni.retiarii.nn.pytorch.api.InputChoice'))
for node_list in ic_nodes:
assert _is_all_equal(map(lambda node: node.operation.parameters['n_candidates'], node_list)) and \
_is_all_equal(map(lambda node: node.operation.parameters['n_chosen'], node_list)), \
'Input choice with the same label must have the same number of candidates.'
mutator = InputChoiceMutator(node_list)
applied_mutators.append(mutator)
vc_nodes = _group_by_label(model.get_nodes_by_type('__torch__.nni.retiarii.nn.pytorch.api.ValueChoice'))
for node_list in vc_nodes:
assert _is_all_equal(map(lambda node: node.operation.parameters['candidates'], node_list)), \
'Value choice with the same label must have the same candidates.'
mutator = ValueChoiceMutator(node_list, node_list[0].operation.parameters['candidates'])
applied_mutators.append(mutator)
pc_nodes = []
for node in model.get_nodes():
for name, choice in node.operation.parameters.items():
if isinstance(choice, ValueChoice):
pc_nodes.append((node, name))
pc_nodes = _group_parameters_by_label(pc_nodes)
for node_list in pc_nodes:
assert _is_all_equal([node.operation.parameters[name].candidates for node, name in node_list]), \
'Value choice with the same label must have the same candidates.'
first_node, first_argname = node_list[0]
mutator = ParameterChoiceMutator(node_list, first_node.operation.parameters[first_argname].candidates)
applied_mutators.append(mutator)
# apply layer choice at last as it will delete some nodes
lc_nodes = _group_by_label(filter(lambda d: d.operation.parameters.get('mutation') == 'layerchoice',
model.get_nodes_by_type('_cell')))
for node_list in lc_nodes:
assert _is_all_equal(map(lambda node: len(node.operation.parameters['candidates']), node_list)), \
'Layer choice with the same label must have the same number of candidates.'
mutator = LayerChoiceMutator(node_list)
applied_mutators.append(mutator)
repeat_nodes = _group_by_label(filter(lambda d: d.operation.parameters.get('mutation') == 'repeat',
model.get_nodes_by_type('_cell')))
for node_list in repeat_nodes:
assert _is_all_equal(map(lambda node: node.operation.parameters['max_depth'], node_list)) and \
_is_all_equal(map(lambda node: node.operation.parameters['min_depth'], node_list)), \
'Repeat with the same label must have the same number of candidates.'
mutator = RepeatMutator(node_list)
applied_mutators.append(mutator)
if applied_mutators:
return applied_mutators
return None
def process_inline_mutation(model: Model) -> Optional[List[Mutator]]:
applied_mutators = []
ic_nodes = _group_by_label(
model.get_nodes_by_type(
'__torch__.nni.retiarii.nn.pytorch.api.InputChoice'))
for node_list in ic_nodes:
assert _is_all_equal(map(lambda node: node.operation.parameters['n_candidates'], node_list)) and \
_is_all_equal(map(lambda node: node.operation.parameters['n_chosen'], node_list)), \
'Input choice with the same label must have the same number of candidates.'
mutator = InputChoiceMutator(node_list)
applied_mutators.append(mutator)
vc_nodes = _group_by_label(
model.get_nodes_by_type(
'__torch__.nni.retiarii.nn.pytorch.api.ValueChoice'))
for node_list in vc_nodes:
assert _is_all_equal(map(lambda node: node.operation.parameters['candidates'], node_list)), \
'Value choice with the same label must have the same candidates.'
mutator = ValueChoiceMutator(
node_list, node_list[0].operation.parameters['candidates'])
applied_mutators.append(mutator)
# `pc_nodes` are arguments of basic units. They can be compositions.
pc_nodes: List[Tuple[Node, str, ValueChoiceX]] = []
for node in model.get_nodes():
# arguments used in operators like Conv2d
# argument `valuechoice` used in generated repeat cell
for name, choice in node.operation.parameters.items():
if isinstance(choice, ValueChoiceX):
# e.g., (conv_node, "out_channels", ValueChoice([1, 3]))
pc_nodes.append((node, name, choice))
# Break `pc_nodes` down to leaf value choices. They should be what we want to sample.
leaf_value_choices: Dict[str, List[Any]] = {}
for _, __, choice in pc_nodes:
for inner_choice in choice.inner_choices():
if inner_choice.label not in leaf_value_choices:
leaf_value_choices[
inner_choice.label] = inner_choice.candidates
else:
assert leaf_value_choices[inner_choice.label] == inner_choice.candidates, \
'Value choice with the same label must have the same candidates, but found ' \
f'{leaf_value_choices[inner_choice.label]} vs. {inner_choice.candidates}'
for label, candidates in leaf_value_choices.items():
applied_mutators.append(ParameterChoiceLeafMutator(candidates, label))
# in the end, add another parameter choice mutator for "real" mutations
if pc_nodes:
applied_mutators.append(
ParameterChoiceMutator([(node, name)
for node, name, _ in pc_nodes]))
# apply layer choice at last as it will delete some nodes
lc_nodes = _group_by_label(
filter(
lambda d: d.operation.parameters.get('mutation') == 'layerchoice',
model.get_nodes_by_type('_cell')))
for node_list in lc_nodes:
assert _is_all_equal(map(lambda node: len(node.operation.parameters['candidates']), node_list)), \
'Layer choice with the same label must have the same number of candidates.'
mutator = LayerChoiceMutator(node_list)
applied_mutators.append(mutator)
repeat_nodes = _group_by_label(
filter(lambda d: d.operation.parameters.get('mutation') == 'repeat',
model.get_nodes_by_type('_cell')))
for node_list in repeat_nodes:
# this check is not completely reliable, because it only checks max and min
assert _is_all_equal(map(lambda node: node.operation.parameters['max_depth'], node_list)) and \
_is_all_equal(map(lambda node: node.operation.parameters['min_depth'], node_list)), \
'Repeat with the same label must have the same candidates.'
mutator = RepeatMutator(node_list)
applied_mutators.append(mutator)
if applied_mutators:
return applied_mutators
return None