def build_frozen_conv6_lottery(args, init_values, mask_values): 
    return SequentialNetwork([
        FreezeConv2D(64, 3, init_values[0], init_values[1], mask_values[0], mask_values[1], kernel_initializer=glorot_normal, padding='same', kernel_regularizer=l2reg(args.l2), name='conv2D_1'),
        Activation('relu'),
        FreezeConv2D(64, 3, init_values[2], init_values[3], mask_values[2], mask_values[3], kernel_initializer=glorot_normal, padding='same', kernel_regularizer=l2reg(args.l2), name='conv2D_2'),
        Activation('relu'),
        MaxPooling2D((2, 2), (2, 2)),
        
        FreezeConv2D(128, 3, init_values[4], init_values[5], mask_values[4], mask_values[5], kernel_initializer=glorot_normal, padding='same', kernel_regularizer=l2reg(args.l2), name='conv2D_3'),
        Activation('relu'),
        FreezeConv2D(128, 3, init_values[6], init_values[7], mask_values[6], mask_values[7], kernel_initializer=glorot_normal, padding='same', kernel_regularizer=l2reg(args.l2), name='conv2D_4'),
        Activation('relu'),
        MaxPooling2D((2, 2), (2, 2)),
        
        FreezeConv2D(256, 3, init_values[8], init_values[9], mask_values[8], mask_values[9], kernel_initializer=glorot_normal, padding='same', kernel_regularizer=l2reg(args.l2), name='conv2D_5'),
        Activation('relu'),
        FreezeConv2D(256, 3, init_values[10], init_values[11], mask_values[10], mask_values[11], kernel_initializer=glorot_normal, padding='same', kernel_regularizer=l2reg(args.l2), name='conv2D_6'),
        Activation('relu'),
        MaxPooling2D((2, 2), (2, 2)),
        Flatten(),
#         Dropout(0.5),
        FreezeDense(256, init_values[12], init_values[13], mask_values[12], mask_values[13], kernel_initializer=glorot_normal, activation=relu, kernel_regularizer=l2reg(args.l2), name='fc_1'),
#         Dropout(0.5),
        FreezeDense(256, init_values[14], init_values[15], mask_values[14], mask_values[15], kernel_initializer=glorot_normal, activation=relu, kernel_regularizer=l2reg(args.l2), name='fc_2'),
#         Dropout(0.5),
        FreezeDense(10, init_values[16], init_values[17], mask_values[16], mask_values[17], kernel_initializer=glorot_normal, activation=None, kernel_regularizer=l2reg(args.l2), name='fc_3')
    ])
def build_conv6_lottery(args): 
    return SequentialNetwork([
        Conv2D(64, 3, kernel_initializer=glorot_normal, padding='same', kernel_regularizer=l2reg(args.l2), name='conv2D_1'),
        Activation('relu'),
        Conv2D(64, 3, kernel_initializer=glorot_normal, padding='same', kernel_regularizer=l2reg(args.l2), name='conv2D_2'),
        Activation('relu'),
        MaxPooling2D((2, 2), (2, 2)),
        Conv2D(128, 3, kernel_initializer=glorot_normal, padding='same', kernel_regularizer=l2reg(args.l2), name='conv2D_3'),
        Activation('relu'),
        Conv2D(128, 3, kernel_initializer=glorot_normal, padding='same', kernel_regularizer=l2reg(args.l2), name='conv2D_4'),
        Activation('relu'),
        MaxPooling2D((2, 2), (2, 2)),
        Conv2D(256, 3, kernel_initializer=glorot_normal, padding='same', kernel_regularizer=l2reg(args.l2), name='conv2D_5'),
        Activation('relu'),
        Conv2D(256, 3, kernel_initializer=glorot_normal, padding='same', kernel_regularizer=l2reg(args.l2), name='conv2D_6'),
        Activation('relu'),
        MaxPooling2D((2, 2), (2, 2)),
        Flatten(),
#         Dropout(0.5),
        Dense(256, kernel_initializer=glorot_normal, activation=relu, kernel_regularizer=l2reg(args.l2), name='fc_1'),
#         Dropout(0.5),
        Dense(256, kernel_initializer=glorot_normal, activation=relu, kernel_regularizer=l2reg(args.l2), name='fc_2'),
#         Dropout(0.5),
        Dense(10, kernel_initializer=glorot_normal, activation=None, kernel_regularizer=l2reg(args.l2), name='fc_3')
    ])
def build_vgg_mini(args):
    return SequentialNetwork([
        Conv2D(64, (3, 3), kernel_initializer=he_normal, padding='same', activation=relu, kernel_regularizer=l2reg(args.l2), name='conv2D_1'),
        MaxPooling2D((2, 2), (2, 2)),
        Conv2D(128, (3, 3), kernel_initializer=he_normal, padding='same', activation=relu, kernel_regularizer=l2reg(args.l2), name='conv2D_2'),
        MaxPooling2D((2, 2), (2, 2)),
        Conv2D(256, (3, 3), kernel_initializer=he_normal, padding='same', activation=relu, kernel_regularizer=l2reg(args.l2), name='conv2D_3'),
        MaxPooling2D((2, 2), (2, 2)),
        Flatten(),
        Dense(512, kernel_initializer=he_normal, activation=relu, kernel_regularizer=l2reg(args.l2), name='fc_1'),
        Dropout(0.5),
        Dense(10, kernel_initializer=he_normal, activation=None, kernel_regularizer=l2reg(args.l2), name='fc_2')
    ])
def build_lenet_conv(args): # ok this is a slightly modified lenet
    return SequentialNetwork([
        Conv2D(20, 5, kernel_initializer=he_normal, padding='same', kernel_regularizer=l2reg(args.l2), name='conv2D_1'),
        # BatchNormalization(momentum=0.0, name='batch_norm_1'),
        Activation('relu'),
        MaxPooling2D((2, 2), (2, 2)),
        Conv2D(40, 5, kernel_initializer=he_normal, padding='same', kernel_regularizer=l2reg(args.l2), name='conv2D_2'),
        # BatchNormalization(momentum=0.0, name='batch_norm_2'),
        Activation('relu'),
        MaxPooling2D((2, 2), (2, 2)),
        Flatten(),
        Dropout(0.25),
        Dense(400, kernel_initializer=he_normal, activation=relu, kernel_regularizer=l2reg(args.l2), name='fc_1'),
        Dropout(0.5),
        Dense(10, kernel_initializer=he_normal, activation=None, kernel_regularizer=l2reg(args.l2), name='fc_2')
    ])
def build_conv2_supermask(args):
    kwargs = {}
    if args.signed_constant:
        kwargs['signed_constant'] = True
        kwargs['const_multiplier'] = args.signed_constant_multiplier
    if args.dynamic_scaling:
        kwargs['dynamic_scaling'] = True

    return SequentialNetwork([
        MaskedConv2D(64,
                     3,
                     kernel_initializer=glorot_normal,
                     sigmoid_bias=args.sigmoid_bias,
                     round_mask=args.round_mask,
                     padding='same',
                     kernel_regularizer=l2reg(args.l2),
                     name='conv2D_1',
                     **kwargs),
        Activation('relu'),
        MaskedConv2D(64,
                     3,
                     kernel_initializer=glorot_normal,
                     sigmoid_bias=args.sigmoid_bias,
                     round_mask=args.round_mask,
                     padding='same',
                     kernel_regularizer=l2reg(args.l2),
                     name='conv2D_2',
                     **kwargs),
        Activation('relu'),
        MaxPooling2D((2, 2), (2, 2)),
        Flatten(),
        MaskedDense(256,
                    kernel_initializer=glorot_normal,
                    sigmoid_bias=args.sigmoid_bias,
                    round_mask=args.round_mask,
                    activation=relu,
                    kernel_regularizer=l2reg(args.l2),
                    name='fc_1',
                    **kwargs),
        MaskedDense(256,
                    kernel_initializer=glorot_normal,
                    sigmoid_bias=args.sigmoid_bias,
                    round_mask=args.round_mask,
                    activation=relu,
                    kernel_regularizer=l2reg(args.l2),
                    name='fc_2',
                    **kwargs),
        MaskedDense(10,
                    kernel_initializer=glorot_normal,
                    sigmoid_bias=args.sigmoid_bias,
                    round_mask=args.round_mask,
                    activation=None,
                    kernel_regularizer=l2reg(args.l2),
                    name='fc_3',
                    **kwargs)
    ])
def build_linknet_2(args):
    layers = conv_bn_relu(32, 3, stride=1, name="block1_conv1")
    for layer in conv_bn_relu(32, 3, stride=1, name="block1_conv2"):
        layers.append(layer)
    layers.append(MaxPooling2D((2, 2), strides=(2, 2), padding="same", name="block1_pool"))
    layers.append(Activation('relu'))
    layers.append(Flatten())
    layers.append(Dense(400, kernel_initializer=he_normal, activation=relu, kernel_regularizer=l2reg(0), name='fc_1'))
    layers.append(Dense(10, kernel_initializer=he_normal, activation=None, kernel_regularizer=l2reg(0), name='fc_2'))
    return SequentialNetwork(layers)
def build_basic_model(args):
    return SequentialNetwork([
        Conv2D(16, 2, padding='same', name='conv2D_1'),
        # BatchNormalization(momentum=0.0, name='batch_norm_1'),
        Activation('relu'),
        MaxPooling2D((2, 2), (2, 2)),
        Conv2D(32, 2, padding='same', name='conv2D_2'),
        # BatchNormalization(momentum=0.0, name='batch_norm_2'),
        Activation('relu'),
        MaxPooling2D((2, 2), (2, 2)),
        Conv2D(64, 2, padding='same', name='conv2D_3'),
        Activation('relu'),
        MaxPooling2D((2,2), (2,2)),
        GlobalAveragePooling2D(),
        Dense(1000, activation=relu),
        Dropout(0.2),
        Dense(1000, activation=relu, name='fc_1'),
        Dropout(0.2),
        Dense(5, activation= None, name='fc_2')
    ])
def build_masked_conv2_lottery(args, mask_values): 
    return SequentialNetwork([
        MaskedConv2D(64, 3, mask_values[0], mask_values[1], kernel_initializer=glorot_normal, padding='same', kernel_regularizer=l2reg(args.l2), name='conv2D_1'),
        Activation('relu'),
        MaskedConv2D(64, 3, mask_values[2], mask_values[3], kernel_initializer=glorot_normal, padding='same', kernel_regularizer=l2reg(args.l2), name='conv2D_2'),
        Activation('relu'),
        MaxPooling2D((2, 2), (2, 2)),
        Flatten(),
        MaskedDense(256, mask_values[4], mask_values[5], kernel_initializer=glorot_normal, activation=relu, kernel_regularizer=l2reg(args.l2), name='fc_1'),
        MaskedDense(256, mask_values[6], mask_values[7], kernel_initializer=glorot_normal, activation=relu, kernel_regularizer=l2reg(args.l2), name='fc_2'),
        MaskedDense(10, mask_values[8], mask_values[9], kernel_initializer=glorot_normal, activation=None, kernel_regularizer=l2reg(args.l2), name='fc_3')
    ])
    def __init__(self, classes=1, dropout=0.5, feature_scale=4):
        super(LinkNet, self).__init__()

        self.conv_bn_relu_1 = self.track_layers(conv_bn_relu(32, 3, stride=1, name="block1_conv1"))
        self.conv_bn_relu_2 = self.track_layers(conv_bn_relu(32, 3, stride=1, name="block1_conv2"))

        self.maxPool = MaxPooling2D((2, 2), strides=(2, 2), padding="same", name="block1_pool")

        layers = [2, 2, 2, 2, 2]
        filters = [64, 128, 256, 512, 32]

        enc1 = self.track_layers(encoder(m=32, n=filters[0], blocks=layers[0], stride=1, name='encoder1'))
        enc2 = self.track_layers(encoder(m=filters[0], n=filters[1], blocks=layers[1], stride=2, name='encoder2'))
        enc3 = self.track_layers(encoder(m=filters[1], n=filters[2], blocks=layers[2], stride=2, name='encoder3'))
        enc4 = self.track_layers(encoder(m=filters[2], n=filters[3], blocks=layers[3], stride=2, name='encoder4'))
        enc5 = self.track_layers(encoder(m=filters[3], n=filters[4], blocks=layers[4], stride=2, name='encoder5'))

        self.decoder = self.track_layer(LinkNetDecoder(enc1, enc2, enc3, enc4, enc5, filters, feature_scale))
        # self.dropout = tfkeras.layers.SpatialDropout2D(dropout)
        self.conv1 = self.track_layer(Conv2D(filters=classes, kernel_size=(1, 1), padding='same', name='prediction'))
        self.act = self.track_layer(Activation('sigmoid', name='mask'))
    def __init__(self,
                 rpn_params,
                 bsamp_params,
                 nms_params,
                 l2=0,
                 im_h=64,
                 im_w=64,
                 coordconv=False,
                 clip=True,
                 filtersame=False):
        super(RegionProposalSampler, self).__init__()
        self.rpn_params = rpn_params
        self.bsamp_params = bsamp_params
        self.nms_params = nms_params
        self.im_h = im_h
        self.im_w = im_w
        self.clip = clip

        _pad = 'same' if filtersame else 'valid'
        _dim = 16 if filtersame else 13
        if coordconv:
            self.l(
                'bottom_conv',
                SequentialNetwork(
                    [
                        AddCoords(x_dim=im_w,
                                  y_dim=im_h,
                                  with_r=False,
                                  skiptile=True),  # (batch, 64, 64, 4 or 5)
                        Conv2D(32, (5, 5),
                               padding=_pad,
                               kernel_initializer=he_normal,
                               kernel_regularizer=l2reg(l2)),
                        ReLu,
                        MaxPooling2D(pool_size=2, strides=2),
                        Conv2D(64, (5, 5),
                               padding=_pad,
                               kernel_initializer=he_normal,
                               kernel_regularizer=l2reg(l2)),
                        ReLu,
                        MaxPooling2D(pool_size=2, strides=2),
                    ],
                    name='bottom_conv'))

            self.l(
                'another_conv',
                SequentialNetwork([
                    AddCoords(
                        x_dim=_dim, y_dim=_dim, with_r=False, skiptile=True),
                    Conv2D(rpn_params.rpn_hidden_dim, (3, 3),
                           padding='same',
                           kernel_initializer=he_normal,
                           kernel_regularizer=l2reg(l2)), ReLu
                ],
                                  name='another_conv'))

            self.l(
                'box_mover',
                SequentialNetwork([
                    Conv2D(rpn_params.rpn_hidden_dim, (3, 3),
                           padding='same',
                           kernel_initializer=he_normal,
                           kernel_regularizer=l2reg(l2)), ReLu,
                    AddCoords(
                        x_dim=_dim, y_dim=_dim, with_r=False, skiptile=True),
                    Conv2D(4 * rpn_params.num_anchors, (1, 1),
                           kernel_initializer=tf.zeros_initializer,
                           bias_initializer=tf.constant_initializer([0.]),
                           kernel_regularizer=l2reg(l2))
                ],
                                  name='box_mover'))  # (13,13,4*k)

        else:
            self.l(
                'bottom_conv',
                SequentialNetwork([
                    Conv2D(32, (5, 5),
                           padding=_pad,
                           kernel_initializer=he_normal,
                           kernel_regularizer=l2reg(l2)),
                    ReLu,
                    MaxPooling2D(pool_size=2, strides=2),
                    Conv2D(64, (5, 5),
                           padding=_pad,
                           kernel_initializer=he_normal,
                           kernel_regularizer=l2reg(l2)),
                    ReLu,
                    MaxPooling2D(pool_size=2, strides=2),
                ],
                                  name='bottom_conv'))

            self.l(
                'another_conv',
                SequentialNetwork([
                    Conv2D(rpn_params.rpn_hidden_dim, (3, 3),
                           padding='same',
                           kernel_initializer=he_normal,
                           kernel_regularizer=l2reg(l2)), ReLu
                ],
                                  name='another_conv'))

            self.l('box_mover',
                   SequentialNetwork([
                       Conv2D(rpn_params.rpn_hidden_dim, (3, 3),
                              padding='same',
                              kernel_initializer=he_normal,
                              kernel_regularizer=l2reg(l2)), ReLu,
                       Conv2D(4 * rpn_params.num_anchors, (1, 1),
                              kernel_initializer=tf.zeros_initializer,
                              bias_initializer=tf.constant_initializer([0.]),
                              kernel_regularizer=l2reg(l2))
                   ],
                                     name='box_mover'))  # (13,13,4*k)

        self.l('box_scorer',
               SequentialNetwork([
                   Conv2D(2 * rpn_params.num_anchors, (1, 1),
                          kernel_initializer=he_normal,
                          kernel_regularizer=l2reg(l2)),
               ],
                                 name='box_scorer'))  # (13,13,2*k)

        return