def _set_cell_ops(edge, C):
edge.data.set('op', [
ops.Identity(),
ops.Zero(stride=1),
ops.ReLUConvBN(C, C, kernel_size=3),
ops.ReLUConvBN(C, C, kernel_size=1),
ops.AvgPool1x1(kernel_size=3, stride=1),
])
def _set_cell_ops(edge, C, stride):
"""
Set the primitives for the bottom level motif where we
have actual ops at the edges.
"""
if isinstance(edge.data.op, list) and all(isinstance(op, Graph) for op in edge.data.op):
return # We are at the edge of an motif
elif isinstance(edge.data.op, ops.Identity):
edge.data.set('op', [
ops.Identity() if stride==1 else ops.FactorizedReduce(C, C),
ops.Zero1x1(stride=stride),
ops.MaxPool1x1(3, stride),
ops.AvgPool1x1(3, stride),
ops.SepConv(C, C, kernel_size=3, stride=stride, padding=1, affine=False),
DepthwiseConv(C, C, kernel_size=3, stride=stride, padding=1, affine=False),
ConvBNReLU(C, C, kernel_size=1),
])
else:
raise ValueError()
def _set_ops(edge, C, stride):
"""
Replace the 'op' at the edges with the ones defined here.
This function is called by the framework for every edge in
the defined scope.
Args:
current_egde_data (EdgeData): The data that currently sits
at the edge.
C (int): convolutional channels
stride (int): stride for the operation
Returns:
EdgeData: the updated EdgeData object.
"""
edge.data.set('op', [
ops.Identity()
if stride == 1 else FactorizedReduce(C, C, stride, affine=False),
ops.Zero(stride=stride),
ops.MaxPool(3, stride),
ops.AvgPool(3, stride),
ops.SepConv(
C, C, kernel_size=3, stride=stride, padding=1, affine=False),
ops.SepConv(
C, C, kernel_size=5, stride=stride, padding=2, affine=False),
ops.DilConv(C,
C,
kernel_size=3,
stride=stride,
padding=2,
dilation=2,
affine=False),
ops.DilConv(C,
C,
kernel_size=5,
stride=stride,
padding=4,
dilation=2,
affine=False),
])
def _set_cell_ops(edge, C):
edge.data.set('op', [
ops.Identity(),
ops.Zero(stride=1),
])
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 _set_cell_ops(current_edge_data, C):
current_edge_data.set('op', [
ops.Identity(),
ops.Zero(stride=1),
])
