def make_model():
set_seeds()
model = SeaNet({
1: (mm.MorphConv2d(1, 32, kernel_size=3, padding=1), 0),
2: (mm.MorphConv2d(32, 32, kernel_size=3, padding=1), 1),
3: (mm.MorphBCRLayer(32, 32, kernel_size=3, padding=1), 2),
4: (nn.MaxPool2d(2), 3),
5: (mm.MorphFlatLinear(6272, 10), 4)
})
model.eval()
return model
def do_flat_insert(model, idx=None):
# !! Can't be inserting ReLUs willy-nilly
if idx is None:
idx2, idx1 = model.random_edge()
else:
idx1, idx2 = idx
print(colstring.blue("do_flat_insert: %d -> %d" % (idx1, idx2)), file=sys.stderr)
old_layer1 = model.get_node(idx1)
old_layer2 = model.get_node(idx2)
layer_factories = [
lambda c: mm.MorphConv2d(c, c, kernel_size=3, padding=1).to_eye(),
lambda c: mm.MorphBatchNorm2d(c).to_eye(),
lambda c: mm.ReLU(),
]
layer_factory = layer_factories[np.random.choice(len(layer_factories))]
def f(model, idx1=idx1, idx2=idx2, layer_factory=layer_factory):
channels = model.forward(layer=idx1).shape[1]
new_layer = layer_factory(channels)
_ = model.modify_edge(idx1, idx2, new_layer)
model.compile()
return model
return f
def make_seablock(in_planes, planes, stride=1, input_dim=32):
if stride == 1:
# simple
return SeaNet(
{
1: (mm.MorphConv2d(in_planes,
planes,
kernel_size=3,
padding=1,
stride=1,
bias=False), 0),
2: (mm.AddLayer(alpha=0.5), [1, 0]),
3: (mm.IdentityLayer(),
2), # To allow layers to be added at the end
},
input_shape=(in_planes, input_dim, input_dim),
tags='simple')
else:
# downsample
return SeaNet(
{
1: (mm.MorphConv2d(in_planes,
planes,
kernel_size=3,
padding=1,
stride=stride,
bias=False), 0),
2: (mm.MorphConv2d(in_planes,
planes,
kernel_size=1,
stride=stride,
bias=False), 0),
3: (mm.AddLayer(alpha=0.5), [2, 1]),
4: (mm.IdentityLayer(),
3), # To allow layers to be added at the end
},
input_shape=(in_planes, input_dim, input_dim),
tags='downsample')
def f(model, idx1=idx1, idx2=idx2):
l1_size = model(layer=idx1).size()
l2_size = model(layer=idx2).size()
merge_node = model.add_skip(idx1, idx2, mm.AddLayer())
if _need_fix_channels:
new_layer = mm.MorphConv2d(l1_size[1], l2_size[1], kernel_size=1, padding=0)
idx1 = model.modify_edge(idx1, merge_node, new_layer)
if _need_fix_spatial:
scale = int(l1_size[-1] / l2_size[-1])
_ = model.modify_edge(idx1, merge_node, mm.MaxPool2d(scale))
model.compile()
return model
def f(model, idx1=idx1, idx2=idx2):
l1_size = model(layer=idx1).size()
l2_size = model(layer=idx2).size()
# Add skip
merge_node = model.add_skip(idx1, idx2, mm.CatLayer())
# If different spatial extent, add max pooling
if l1_size[-1] != l2_size[-1]:
if l1_size[-1] > l2_size[-1]:
scale = int(l1_size[-1] / l2_size[-1])
maxpool_node = model.modify_edge(idx1, merge_node, mm.MaxPool2d(scale))
else:
scale = int(l2_size[-1] / l1_size[-1])
maxpool_node = model.modify_edge(idx2, merge_node, mm.MaxPool2d(scale))
# 1x1 conv to fix output dim
out_size = model(layer=merge_node).size()
new_layer = mm.MorphConv2d(out_size[1], l2_size[1], kernel_size=1, padding=0)
new_layer = new_layer.to_eye()
model.modify_edge(merge_node, None, new_layer)
model.compile()
return model