Exemplo n.º 1
0
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
Exemplo n.º 2
0
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
Exemplo n.º 3
0
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')
Exemplo n.º 4
0
 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
Exemplo n.º 5
0
 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