def _set_makrograph_ops(self, channel_map_from, channel_map_to, reduction_cell_indices, max_index, affine=True):
# pre-processing
# In darts there is a hardcoded multiplier of 3 for the output of the stem
stem_multiplier = 3
self.edges[1, 2].set('op', ops.Stem(self.channels[0] * stem_multiplier))
# edges connecting cells
for u, v, data in sorted(self.edges(data=True)):
if u > 1 and v < max_index:
C_in = self.channels[channel_map_from[u]]
C_out = self.channels[channel_map_to[v]]
if C_in == C_out:
C_in = C_in * stem_multiplier if u == 2 else C_in * self.num_in_edges # handle Stem
if v in reduction_cell_indices:
C_out *= 2
data.set('op', ops.ReLUConvBN(C_in, C_out, kernel_size=1, affine=affine))
else:
data.set('op', FactorizedReduce(C_in * self.num_in_edges, C_out, affine=affine))
# post-processing
_, _, data = sorted(self.edges(data=True))[-1]
data.set('op', ops.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(self.channels[-1] * self.num_in_edges, self.num_classes))
)
def prepare_evaluation(self):
"""
In DARTS the evaluation model has 32 channels after the Stem
and 3 normal cells at each stage.
"""
# this is called after the optimizer has discretized the graph
self._expand()
# Operations at the edges
self.channels = [36, 72, 144]
reduction_cell_indices = [9, 16]
channel_map_from, channel_map_to = channel_maps(reduction_cell_indices, max_index=23)
self._set_makrograph_ops(channel_map_from, channel_map_to, reduction_cell_indices, max_index=23, affine=True)
# Taken from DARTS implementation
# assuming input size 8x8
self.edges[22, 23].set('op', ops.Sequential(
nn.ReLU(inplace=True),
nn.AvgPool2d(5, stride=3, padding=0, count_include_pad=False), # image size = 2 x 2
nn.Conv2d(self.channels[-1] * self.num_in_edges, 128, 1, bias=False),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
nn.Conv2d(128, 768, 2, bias=False),
nn.BatchNorm2d(768),
nn.ReLU(inplace=True),
nn.Flatten(),
nn.Linear(768, self.num_classes))
)
self.update_edges(
update_func=_double_channels,
scope=self.OPTIMIZER_SCOPE,
private_edge_data=True
)
def __init__(self):
super().__init__()
# Define the motifs (6 level-2 motifs)
level2_motifs = []
for j in range(6):
motif = Graph()
motif.name = "motif{}".format(j)
motif.add_nodes_from([i for i in range(1, 5)])
motif.add_edges_densly()
level2_motifs.append(motif)
# cell (= one level-3 motif)
cell = Graph()
cell.name = "cell"
cell.add_nodes_from([i for i in range(1, 6)])
cell.add_edges_densly()
cells = []
channels = [16, 32, 64]
for scope, c in zip(self.OPTIMIZER_SCOPE, channels):
cell_i = cell.copy().set_scope(scope)
cell_i.update_edges(
update_func=lambda edge: _set_motifs(edge, motifs=level2_motifs, c=c),
private_edge_data=True
)
cell_i.set_scope(scope)
# set the level 1 motifs (i.e. primitives)
cell_i.update_edges(
update_func=lambda edge: _set_cell_ops(edge, c, stride=1),
scope=[scope],
private_edge_data=True
)
cells.append(cell_i)
self.name = "makrograph"
self.add_nodes_from([i for i in range(1, 9)])
self.add_edges_from([(i, i+1) for i in range(1, 8)])
self.edges[1, 2].set('op', ops.Stem(16))
self.edges[2, 3].set('op', cells[0])
self.edges[3, 4].set('op', ops.SepConv(16, 32, kernel_size=3, stride=2, padding=1))
self.edges[4, 5].set('op', cells[1])
self.edges[5, 6].set('op', ops.SepConv(32, 64, kernel_size=3, stride=2, padding=1))
self.edges[6, 7].set('op', cells[2])
self.edges[7, 8].set('op', ops.Sequential(
ops.SepConv(64, 64, kernel_size=3, stride=1, padding=1),
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(channels[-1], 10))
)
def prepare_evaluation(self):
"""
The evaluation model has N=2 cells at each stage and a sepconv with stride 1
between them. Initial channels = 64, trained 512 epochs. Learning rate 0.1
reduced by 10x after 40K, 60K, and 70K steps.
"""
# this is called after the optimizer has discretized the graph
cells = [self.edges[2, 3].op, self.edges[4, 5].op, self.edges[6, 7].op]
self._expand()
channels = [64, 128, 256]
factor = 4
for cell, c in zip(cells, channels):
for _, _, data in cell.edges.data():
data.op.update_nodes(
lambda node, in_edges, out_edges: _set_comb_op_channels(node, in_edges, out_edges, c=c),
single_instances=False
)
self.edges[1, 2].set('op', ops.Stem(channels[0]))
self.edges[2, 3].set('op', cells[0].copy())
self.edges[3, 4].set('op', ops.SepConv(channels[0], channels[0], kernel_size=3, stride=1, padding=1))
self.edges[4, 5].set('op', cells[0].copy())
self.edges[5, 6].set('op', ops.SepConv(channels[0], channels[1], kernel_size=3, stride=2, padding=1))
self.edges[6, 7].set('op', cells[1].copy())
self.edges[7, 8].set('op', ops.SepConv(channels[1], channels[1], kernel_size=3, stride=1, padding=1))
self.edges[8, 9].set('op', cells[1].copy())
self.edges[9, 10].set('op', ops.SepConv(channels[1], channels[2], kernel_size=3, stride=2, padding=1))
self.edges[10, 11].set('op', cells[2].copy())
self.edges[11, 12].set('op', ops.SepConv(channels[2], channels[2], kernel_size=3, stride=1, padding=1))
self.edges[12, 13].set('op', cells[2].copy())
self.edges[13, 14].set('op', ops.Sequential(
ops.SepConv(channels[-1], channels[-1], kernel_size=3, stride=1, padding=1),
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(channels[-1], 10))
)
self.update_edges(
update_func=lambda edge: _increase_channels(edge, factor),
scope=self.OPTIMIZER_SCOPE,
private_edge_data=True
)
def __init__(self):
logger.info("This is not a search space but the final model found by "
"Liu et al. For the search space use `HierarchicalSearchSpace()`")
# skip the init of HierarchicalSearchSpace
Graph.__init__(self)
separable_3x3 = {
'op': ops.SepConv,
'kernel_size': 3,
'stride': 1,
'padding': 1
}
depthwise_3x3 = {
'op': DepthwiseConv,
'kernel_size': 3,
'stride': 1,
'padding': 1
}
conv_1x1 = {
'op': ConvBNReLU,
'kernel_size': 1,
}
motif = Graph()
motif.add_nodes_from([1, 2, 3, 4])
motif.add_edges_densly()
# Motif 1
motif1 = motif.clone()
motif1.name = "motif1"
motif1.edges[1, 2].set('op', ops.MaxPool1x1(kernel_size=3, stride=1))
motif1.edges[1, 3].update(conv_1x1)
motif1.edges[1, 4].update(separable_3x3)
motif1.edges[2, 3].update(depthwise_3x3)
motif1.edges[2, 4].update(conv_1x1)
motif1.edges[3, 4].set('op', ops.MaxPool1x1(kernel_size=3, stride=1))
# Motif 2
motif2 = motif.clone()
motif2.name = "motif2"
motif2.edges[1, 2].set('op', ops.MaxPool1x1(kernel_size=3, stride=1))
motif2.edges[1, 3].update(depthwise_3x3)
motif2.edges[1, 4].update(conv_1x1)
motif2.edges[2, 3].update(depthwise_3x3)
motif2.edges[2, 4].update(separable_3x3)
motif2.edges[3, 4].update(separable_3x3)
# Motif 3
motif3 = motif.clone()
motif3.name = "motif3"
motif3.edges[1, 2].update(separable_3x3)
motif3.edges[1, 3].update(depthwise_3x3)
motif3.edges[1, 4].update(separable_3x3)
motif3.edges[2, 3].update(separable_3x3)
motif3.edges[2, 4].set('op', ops.Identity()) # assuming no label is identiy
motif3.edges[3, 4].set('op', ops.AvgPool1x1(kernel_size=3, stride=1))
# Motif 4
motif4 = motif.clone()
motif4.name = "motif4"
motif4.edges[1, 2].set('op', ops.AvgPool1x1(kernel_size=3, stride=1))
motif4.edges[1, 3].update(separable_3x3)
motif4.edges[1, 4].set('op', ops.Identity())
motif4.edges[2, 3].update(separable_3x3)
motif4.edges[2, 4].set('op', ops.MaxPool1x1(kernel_size=3, stride=1))
motif4.edges[3, 4].set('op', ops.AvgPool1x1(kernel_size=3, stride=1))
# Motif 5
motif5 = motif.clone()
motif5.name = "motif5"
motif5.edges[1, 2].set('op', ops.Identity()) # Unclear in paper
motif5.edges[1, 3].update(separable_3x3)
motif5.edges[1, 4].update(separable_3x3)
motif5.edges[2, 3].set('op', ops.Identity())
motif5.edges[2, 4].set('op', ops.Identity())
motif5.edges[3, 4].update(separable_3x3)
# Motif 6
motif6 = motif.clone()
motif6.name = "motif6"
motif6.edges[1, 2].update(depthwise_3x3)
motif6.edges[1, 3].set('op', ops.MaxPool1x1(kernel_size=3, stride=1))
motif6.edges[1, 4].update(separable_3x3)
motif6.edges[2, 3].set('op', ops.MaxPool1x1(kernel_size=3, stride=1))
motif6.edges[2, 4].set('op', ops.Identity())
motif6.edges[3, 4].set('op', ops.AvgPool1x1(kernel_size=3, stride=1))
# Cell
cell = Graph("cell")
cell.add_nodes_from([1, 2, 3, 4, 5])
cell.add_edges_densly()
cell.edges[1, 2].set('op', motif5.copy())
cell.edges[1, 3].set('op', motif3.copy())
cell.edges[1, 4].set('op', motif3.copy())
cell.edges[1, 5].set('op', motif4.copy())
cell.edges[2, 3].set('op', motif3.copy())
cell.edges[2, 4].set('op', motif1.copy())
cell.edges[2, 5].set('op', motif5.copy())
cell.edges[3, 4].set('op', motif3.copy())
cell.edges[3, 5].set('op', motif5.copy())
cell.edges[4, 5].set('op', motif5.copy())
cells = []
channels = [16, 32, 64]
for scope, c in zip(self.OPTIMIZER_SCOPE, channels):
cell_i = cell.copy().set_scope(scope)
# Specify channels
cell_i.update_edges(
update_func=lambda edge: edge.data.update({'C_in': c, 'C_out': c}),
private_edge_data=True
)
cells.append(cell_i)
self.name = "makrograph"
self.add_nodes_from([i for i in range(1, 9)])
self.add_edges_from([(i, i+1) for i in range(1, 8)])
self.edges[1, 2].set('op', ops.Stem(16))
self.edges[2, 3].set('op', cells[0])
self.edges[3, 4].set('op', ops.SepConv(16, 32, kernel_size=3, stride=2, padding=1))
self.edges[4, 5].set('op', cells[1])
self.edges[5, 6].set('op', ops.SepConv(32, 64, kernel_size=3, stride=2, padding=1))
self.edges[6, 7].set('op', cells[2])
self.edges[7, 8].set('op', ops.Sequential(
ops.SepConv(64, 64, kernel_size=3, stride=1, padding=1),
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(channels[-1], 10))
)
self.compile()
def __init__(self):
super().__init__()
self.num_classes = self.NUM_CLASSES if hasattr(self, 'NUM_CLASSES') else 10
#
# Cell definition
#
cell = Graph()
cell.name = "cell" # Use the same name for all cells with shared attributes
# Input node
cell.add_node(1)
# Intermediate nodes
cell.add_node(2)
cell.add_node(3)
# Output node
cell.add_node(4)
# Edges
cell.add_edges_densly()
#
# Makrograph definition
#
self.name = "makrograph"
# Cell is on the edges
# 1-2: Preprocessing
# 2-3, ..., 6-7: cells stage 1
# 7-8: residual block stride 2
# 8-9, ..., 12-13: cells stage 2
# 13-14: residual block stride 2
# 14-15, ..., 18-19: cells stage 3
# 19-20: post-processing
total_num_nodes = 20
self.add_nodes_from(range(1, total_num_nodes+1))
self.add_edges_from([(i, i+1) for i in range(1, total_num_nodes)])
channels = [16, 32, 64]
#
# operations at the edges
#
# preprocessing
self.edges[1, 2].set('op', ops.Stem(channels[0]))
# stage 1
for i in range(2, 7):
self.edges[i, i+1].set('op', cell.copy().set_scope('stage_1'))
# stage 2
self.edges[7, 8].set('op', ResNetBasicblock(C_in=channels[0], C_out=channels[1], stride=2))
for i in range(8, 13):
self.edges[i, i+1].set('op', cell.copy().set_scope('stage_2'))
# stage 3
self.edges[13, 14].set('op', ResNetBasicblock(C_in=channels[1], C_out=channels[2], stride=2))
for i in range(14, 19):
self.edges[i, i+1].set('op', cell.copy().set_scope('stage_3'))
# post-processing
self.edges[19, 20].set('op', ops.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(channels[-1], self.num_classes)
))
# set the ops at the cells (channel dependent)
for c, scope in zip(channels, self.OPTIMIZER_SCOPE):
self.update_edges(
update_func=lambda edge: _set_cell_ops(edge, C=c),
scope=scope,
private_edge_data=True
)
def __init__(self,
classes: int = 10,
intermediate_nodes: int = 2,
cells_per_stage: int = 1,
channels: list = [16, 32]):
"""
Initializes the simple cell search space.
Args:
classes (int): Number of classes. Default: 10.
intermediate_nodes (int): Number of intermediate nodes for normal and
reduction cells. Default: 2.
cells_per_stage (int): Number of normal cells at each stage. Default: 1.
channels (list): Channels for each stage. Must have len 2. Default: [16, 32]
"""
assert len(channels) == len(self.OPTIMIZER_SCOPE), \
"Expecting a channel for each scope. Expected {}, got {}.".format(len(self.OPTIMIZER_SCOPE), len(channels))
super().__init__()
# Cell definition
normal_cell = Graph(
name="normal_cell"
) # Use the same name for all cells with shared attributes
# Nodes
out_node_idx = intermediate_nodes + 3
normal_cell.add_nodes_from(range(1, out_node_idx + 1))
# Edges
normal_cell.add_edges_from([(1, i) for i in range(3, out_node_idx)
]) # input 1
normal_cell.add_edges_from([(2, i) for i in range(3, out_node_idx)
]) # input 2
normal_cell.add_edges_from([(u, v)
for u, v in normal_cell.get_dense_edges()
if u not in [1, 2] and v != out_node_idx])
# Edges connecting to the output are always the identity
normal_cell.add_edges_from([(i, out_node_idx, EdgeData().finalize())
for i in range(3, out_node_idx)]) # output
# set the parameters for the ops at all edges (that are not final)
for k, v in edge_attributes.items():
normal_cell.set_at_edges(k, v)
# Reduction cell
reduction_cell = deepcopy(normal_cell)
reduction_cell.name = "reduction_cell"
reduction_cell.update_edges(update_func=lambda edge: edge.data.set(
'stride', 2) if edge.head in [1, 2] else None)
# Makrograph definition
self.name = "makrograph"
self.add_node(1) # input node
self.add_node(2) # preprocessing
self.add_edge(1, 2) # pre-processing (stem)
j = 3 # index of next node to add
for scope in self.OPTIMIZER_SCOPE:
# reduction cell (beginning of each stage but first)
if j > 3:
input = [j - 2, j - 1]
self.add_node(j,
subgraph=reduction_cell.copy().set_scope(
scope).set_input(input))
self.add_edges_from([(input[0], j), (input[1], j)])
j += 1
# normal cells
for _ in range(cells_per_stage):
# single (copied) input if first node or every normal node after reduction cell
input = [
j - 1, j - 1
] if j == 3 or (j - 3) % (cells_per_stage + 1) == 0 else [
j - 2, j - 1
]
self.add_node(j,
subgraph=normal_cell.copy().set_scope(
scope).set_input(input))
self.add_edges_from([(input[0], j), (input[1], j)])
j += 1
self.add_node(j) # output
self.add_edge(j - 1, j) # post-processing (pooling, classifier)
# Compile the ops
self.edges[1, 2].set(
'op', ops.Stem(channels[0])
) # we can also set a compiled op. Will be ignored by compile()
def set_channels(edge, C):
C_in = C if edge.data.stride == 1 else C // 2
edge.data.set('C_in', C_in)
edge.data.set('C_out', C)
for scope, c in zip(self.OPTIMIZER_SCOPE, channels):
self.update_edges(lambda edge: set_channels(edge, c),
scope,
private_edge_data=True)
self.edges[j - 1, j].set(
'op',
ops.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten(),
nn.Linear(channels[-1], classes)))
self.compile()
# Combining operations are currently not considered by compile()
def set_comb_op(node, in_edges, out_edges, C):
if node[0] == out_node_idx:
node[1]['comb_op'] = ops.Concat1x1(
num_in_edges=intermediate_nodes, C_out=C)
for scope, c in zip(self.OPTIMIZER_SCOPE, channels):
self.update_nodes(update_func=lambda node, in_edges, out_edges:
set_comb_op(node, in_edges, out_edges, c),
scope=scope,
single_instances=False)