def AddLeNetModel(model, data):
    '''
    This part is the standard LeNet model: from data to the softmax prediction.

    For each convolutional layer we specify dim_in - number of input channels
    and dim_out - number or output channels. Also each Conv and MaxPool layer changes the
    image size. For example, kernel of size 5 reduces each side of an image by 4.

    While when we have kernel and stride sizes equal 2 in a MaxPool layer, it divides
    each side in half.
    '''
    # Image size: 28 x 28 -> 24 x 24
    conv1 = brew.conv(model, data, 'conv1', dim_in=1, dim_out=20, kernel=5)
    # Image size: 24 x 24 -> 12 x 12
    pool1 = brew.max_pool(model, conv1, 'pool1', kernel=2, stride=2)
    # Image size: 12 x 12 -> 8 x 8
    conv2 = brew.conv(model, pool1, 'conv2', dim_in=20, dim_out=50, kernel=5)
    # Image size: 8 x 8 -> 4 x 4
    pool2 = brew.max_pool(model, conv2, 'pool2', kernel=2, stride=2)
    # 50 * 4 * 4 stands for dim_out from previous layer multiplied by the image size
    fc3 = brew.fc(model, pool2, 'fc3', dim_in=50 * 4 * 4, dim_out=500)
    fc3 = brew.relu(model, fc3, fc3)
    pred = brew.fc(model, fc3, 'pred', 500, 10)
    softmax = brew.softmax(model, pred, 'softmax')
    return softmax 

def create_model(model_builder, model, enable_tensor_core, float16_compute, loss_scale=1.0):
    """Creates one model replica.

    :param obj model_builder: A model instance that contains `forward_pass_builder` method.
    :param model: Caffe2's model helper class instances.
    :type model: :py:class:`caffe2.python.model_helper.ModelHelper`
    :param bool enable_tensor_core: If true, Volta's tensor core ops are enabled.
    :param float loss_scale: Scale loss for multi-GPU training.
    :return: Head nodes (softmax or loss depending on phase)
    """
    initializer = (pFP16Initializer if model_builder.dtype == 'float16' else Initializer)
    with brew.arg_scope([brew.conv, brew.fc],
                        WeightInitializer=initializer,
                        BiasInitializer=initializer,
                        enable_tensor_core=enable_tensor_core,
                        float16_compute=float16_compute):
        outputs = model_builder.forward_pass_builder(model, loss_scale=loss_scale)
    return outputs 

def create_model(m, device_opts) :
    with core.DeviceScope(device_opts):

        conv1 = brew.conv(m, 'data', 'conv1', dim_in=1, dim_out=20, kernel=5)
        pool1 = brew.max_pool(m, conv1, 'pool1', kernel=2, stride=2)
        conv2 = brew.conv(m, pool1, 'conv2', dim_in=20, dim_out=50, kernel=5)
        pool2 = brew.max_pool(m, conv2, 'pool2', kernel=2, stride=2)
        fc3 = brew.fc(m, pool2, 'fc3', dim_in=50 * 4 * 4, dim_out=500)
        fc3 = brew.relu(m, fc3, fc3)
        pred = brew.fc(m, fc3, 'pred', 500, 2)
        softmax = brew.softmax(m, pred, 'softmax')
        m.net.AddExternalOutput(softmax)
        return softmax

# add loss and optimizer 

def AddLeNetModel(model, data):
    '''
    This part is the standard LeNet model: from data to the softmax prediction.

    For each convolutional layer we specify dim_in - number of input channels
    and dim_out - number or output channels. Also each Conv and MaxPool layer changes the
    image size. For example, kernel of size 5 reduces each side of an image by 4.

    While when we have kernel and stride sizes equal 2 in a MaxPool layer, it divides
    each side in half.
    '''
    # Image size: 28 x 28 -> 24 x 24
    conv1 = brew.conv(model, data, 'conv1', dim_in=1, dim_out=20, kernel=5)
    # Image size: 24 x 24 -> 12 x 12
    pool1 = brew.max_pool(model, conv1, 'pool1', kernel=2, stride=2)
    # Image size: 12 x 12 -> 8 x 8
    conv2 = brew.conv(model, pool1, 'conv2', dim_in=20, dim_out=100, kernel=5)
    # Image size: 8 x 8 -> 4 x 4
    pool2 = brew.max_pool(model, conv2, 'pool2', kernel=2, stride=2)
    # 50 * 4 * 4 stands for dim_out from previous layer multiplied by the
    # image size
    fc3 = brew.fc(model, pool2, 'fc3', dim_in=100 * 4 * 4, dim_out=500)
    relu = brew.relu(model, fc3, fc3)
    pred = brew.fc(model, relu, 'pred', 500, 10)
    softmax = brew.softmax(model, pred, 'softmax')
    return softmax 

def test_simple_model(self):
        model = model_helper.ModelHelper(name="mnist")
        # how come those inputs don't break the forward pass =.=a
        workspace.FeedBlob("data", np.random.randn(1, 3, 64, 64).astype(np.float32))
        workspace.FeedBlob("label", np.random.randn(1, 1000).astype(np.int))

        with core.NameScope("conv1"):
            conv1 = brew.conv(model, "data", 'conv1', dim_in=1, dim_out=20, kernel=5)
            # Image size: 24 x 24 -> 12 x 12
            pool1 = brew.max_pool(model, conv1, 'pool1', kernel=2, stride=2)
            # Image size: 12 x 12 -> 8 x 8
            conv2 = brew.conv(model, pool1, 'conv2', dim_in=20, dim_out=100, kernel=5)
            # Image size: 8 x 8 -> 4 x 4
            pool2 = brew.max_pool(model, conv2, 'pool2', kernel=2, stride=2)
        with core.NameScope("classifier"):
            # 50 * 4 * 4 stands for dim_out from previous layer multiplied by the image size
            fc3 = brew.fc(model, pool2, 'fc3', dim_in=100 * 4 * 4, dim_out=500)
            relu = brew.relu(model, fc3, fc3)
            pred = brew.fc(model, relu, 'pred', 500, 10)
            softmax = brew.softmax(model, pred, 'softmax')
            xent = model.LabelCrossEntropy([softmax, "label"], 'xent')
            # compute the expected loss
            loss = model.AveragedLoss(xent, "loss")
        model.net.RunAllOnMKL()
        model.param_init_net.RunAllOnMKL()
        model.AddGradientOperators([loss], skip=1)
        blob_name_tracker = {}
        graph = tb.model_to_graph_def(
            model,
            blob_name_tracker=blob_name_tracker,
            shapes={},
            show_simplified=False,
        )

        compare_proto(graph, self) 

def conv_factory(self, model, v, num_in_channels, num_filters, kernel,
                     stride=1, pad=0, relu=True, name='conv'):
        """Standard convolution block: Conv -> BatchNorm -> Activation
        """
        if isinstance(pad, int):
            pad_t = pad_b = pad_l = pad_r = pad
        elif isinstance(pad, list) or isinstance(pad, tuple):
            if len(pad) == 2:
                pad_t = pad_b = pad[0]
                pad_l = pad_r = pad[1]
            elif len(pad) == 4:
                pad_t = pad[0]
                pad_b = pad[1]
                pad_l = pad[2]
                pad_r = pad[3]
            else:
                assert False, "Invalid length of pad array. Expecting 2 or 4 but have: " + str(pad)
        else:
            assert False, "Invalid type of padding: " + str(pad)

        v = brew.conv(model, v, name + '_conv', num_in_channels, num_filters,
                      kernel=kernel, pad_t=pad_t, pad_l=pad_l, pad_b=pad_b,
                      pad_r=pad_r, stride=stride)
        v = brew.spatial_bn(model, v, name+'_bn', num_filters, eps=2e-5,
                            momentum=0.9, is_test=(self.phase == 'inference'))
        if relu is True:
            v = brew.relu(model, v, name + '_relu')
        return v 

def forward_pass_builder(self, model, loss_scale=1.0):
        """
            This function adds the operators, layers to the network. It should return a list
            of loss-blobs that are used for computing the loss gradient. This function is
            also passed an internally calculated loss_scale parameter that is used to scale
            your loss to normalize for the number of GPUs. Signature: function(model, loss_scale)
        """
        is_inference = self.phase == 'inference'

        v = 'data'

        v = brew.conv(model, v, 'conv1', 3, 64, kernel=11, stride=4)
        v = brew.relu(model, v, 'relu1')
        v = brew.max_pool(model, v, 'pool1', kernel=3, stride=2)

        v = brew.conv(model, v, 'conv2', 64, 192, kernel=5, pad=2, group=1)
        v = brew.relu(model, v, 'relu2')
        v = brew.max_pool(model, v, 'pool2', kernel=3, stride=2)

        v = brew.conv(model, v, 'conv3', 192, 384, kernel=3, pad=1)
        v = brew.relu(model, v, 'relu3')

        v = brew.conv(model, v, 'conv4', 384, 256, kernel=3, pad=1, group=1)
        v = brew.relu(model, v, 'relu4')

        v = brew.conv(model, v, 'conv5', 256, 256, kernel=3, pad=1, group=1)
        v = brew.relu(model, v, 'relu5')
        v = brew.max_pool(model, v, 'pool5', kernel=3, stride=2)

        v = brew.fc(model, v, 'fc6', dim_in=9216, dim_out=4096)
        v = brew.relu(model, v, 'relu6')
        v = brew.dropout(model, v, 'drop6', ratio=0.5, is_test=is_inference)

        v = brew.fc(model, v, 'fc7', dim_in=4096, dim_out=4096)
        v = brew.relu(model, v, 'relu7')
        v = brew.dropout(model, v, 'drop7', ratio=0.5, is_test=is_inference)

        return self.add_head_nodes(model, v, 4096, 'fc8', loss_scale=loss_scale) 

def forward_pass_builder(self, model, loss_scale=1.0):
        """
            This function adds the operators, layers to the network. It should return
            a list of loss-blobs that are used for computing the loss gradient. This
            function is also passed an internally calculated loss_scale parameter that
            is used to scale your loss to normalize for the number of GPUs.
            Signature: function(model, loss_scale)
        """
        is_inference = self.phase == 'inference'
        layers, filters = VGG.specs[self.__model]['specs']
        v = 'data'
        dim_in = self.input_shape[0]
        for i, num in enumerate(layers):
            for j in range(num):
                v = brew.conv(model, v, 'conv%d_%d' % (i+1, j+1), dim_in, filters[i], kernel=3, pad=1)
                v = brew.relu(model, v, 'relu%d_%d' % (i+1, j+1))
                dim_in = filters[i]
            v = brew.max_pool(model, v, 'pool%d' % (i+1), kernel=2, stride=2)

        dim_in = 25088 # 512 * 7 * 7 (output tensor of previous max pool layer)
        for i in range(2):
            v = brew.fc(model, v, 'fc%d' % (6+i), dim_in=dim_in, dim_out=4096)
            v = brew.relu(model, v, 'relu%d' % (6+i))
            v = brew.dropout(model, v, 'drop%d' % (6+i), ratio=0.5, is_test=is_inference)
            dim_in = 4096

        return self.add_head_nodes(model, v, 4096, 'fc8', loss_scale=loss_scale) 

def conv(self, model, name, inputs, input_depth, num_filters, kernel, stride,
             pad, is_inference):
        # Check padding
        if isinstance(pad, int):
            pad_t = pad_b = pad_l = pad_r = pad
        elif isinstance(pad, list) or isinstance(pad, tuple):
            if len(pad) == 2:
                pad_t = pad_b = pad[0]
                pad_l = pad_r = pad[1]
            elif len(pad) == 4:
                pad_t = pad[0]
                pad_b = pad[1]
                pad_l = pad[2]
                pad_r = pad[3]
            else:
                assert False, "Invalid length of pad array. Expecting 2 or 4 but have: " + str(pad)
        else:
            assert False, "Invalid type of padding: " + str(pad)
        # Check kernel
        if isinstance(kernel, int):
            kernel = [kernel, kernel]
        elif isinstance(kernel, tuple) or isinstance(kernel, list):
            assert len(kernel) == 2, "Kernel must have length 2"
            kernel = [kernel[0], kernel[1]]
        else:
            assert False, "Invalid type of kerne;: " + str(kernel)
        #
        self.counts[name] += 1
        name = name + str(self.counts[name]-1)
        #
        v = brew.conv(model, inputs, name + '_conv', input_depth, num_filters,
                      kernel=kernel, stride=stride,
                      pad_t=pad_t, pad_l=pad_l, pad_b=pad_b, pad_r=pad_r,
                      no_bias=True)
        v = brew.spatial_bn(model, v, name+'_bn', num_filters, eps=2e-5,
                            momentum=0.9, is_test=is_inference)
        v = brew.relu(model, v, name+'_relu')
        return v 

def module_a(self, model, inputs, input_depth, n, is_inference):
        branhes = [
            [('conv', 64, 1, 1, 0)],
            [('conv', 48, 1, 1, 0), ('conv', 64, 5, 1, 2)],
            [('conv', 64, 1, 1, 0), ('conv', 96, 3, 1, 1), ('conv', 96, 3, 1, 1)],
            [('avg', 3, 1, 1), ('conv', n, 1, 1, 0)]
        ]
        return self.inception_module(model, 'inception_a', inputs, input_depth,
                                     branhes, is_inference) 

def module_b(self, model, inputs, input_depth, is_inference):
        branches = [
            [('conv', 384, 3, 2, 0)],
            [('conv', 64, 1, 1, 0), ('conv', 96, 3, 1, 1), ('conv', 96, 3, 2, 0)],
            [('max', 3, 2, 0)]
        ]
        return self.inception_module(model, 'inception_b', inputs, input_depth,
                                     branches, is_inference) 

def module_c(self, model, inputs, input_depth, n, is_inference):
        branches = [
            [('conv', 192, 1, 1, 0)],
            [('conv', n, 1, 1, 0), ('conv', n, (1,7), 1, (0,3)), ('conv', 192, (7,1), 1, (3,0))],
            [('conv', n, 1, 1, 0), ('conv', n, (7,1), 1, (3,0)), ('conv', n, (1,7), 1, (0,3)),
             ('conv', n, (7,1), 1, (3,0)), ('conv', 192, (1,7), 1, (0,3))],
            [('avg', 3, 1, 1), ('conv', 192, 1, 1, 0)]
        ]
        return self.inception_module(model, 'inception_c', inputs, input_depth,
                                     branches, is_inference) 

def module_e(self, model, inputs, input_depth, pooltype, is_inference):
        branches = [
            [('conv', 320, 1, 1, 0)],
            [('conv', 384, 1, 1, 0), ('conv', 384, (1,3), 1, (0,1))],
            [('share',),             ('conv', 384, (3,1), 1, (1,0))],
            [('conv', 448, 1, 1, 0), ('conv', 384, 3, 1, 1), ('conv', 384, (1,3), 1, (0,1))],
            [('share',),             ('share',),             ('conv', 384, (3,1), 1, (1,0))],
            [(pooltype, 3, 1, 1), ('conv', 192, 1, 1, 0)]
        ]
        return self.inception_module(model, 'inception_e', inputs, input_depth,
                                     branches, is_inference) 

def forward_pass_builder(self, model, loss_scale=1.0):
        """
            This function adds the operators, layers to the network. It should return a list
            of loss-blobs that are used for computing the loss gradient. This function is
            also passed an internally calculated loss_scale parameter that is used to scale
            your loss to normalize for the number of GPUs. Signature: function(model, loss_scale)
        """
        self.counts = defaultdict(lambda: 0)
        is_inference = self.phase == 'inference'

        v = 'data'

        # Input conv modules
        v = self.conv(model, 'conv', v, input_depth=3, num_filters=32, kernel=3, stride=2, pad=0, is_inference=is_inference)
        v = self.conv(model, 'conv', v, input_depth=32, num_filters=32, kernel=3, stride=1, pad=0, is_inference=is_inference)
        v = self.conv(model, 'conv', v, input_depth=32, num_filters=64, kernel=3, stride=1, pad=1, is_inference=is_inference)
        v = brew.max_pool(model, v, blob_out='pool1', kernel=3, stride=2, pad=0)
        v = self.conv(model, 'conv', v, input_depth=64, num_filters=80, kernel=1, stride=1, pad=0, is_inference=is_inference)
        v = self.conv(model, 'conv', v, input_depth=80, num_filters=192, kernel=3, stride=1, pad=0, is_inference=is_inference)
        v = brew.max_pool(model, v, blob_out='pool2', kernel=3, stride=2, pad=0)
        # Three Type A inception modules
        v = self.module_a(model, inputs=v, input_depth=192, n=32, is_inference=is_inference)
        v = self.module_a(model, inputs=v, input_depth=256, n=64, is_inference=is_inference)
        v = self.module_a(model, inputs=v, input_depth=288, n=64, is_inference=is_inference)
        # One Type B inception module
        v = self.module_b(model, inputs=v, input_depth=288, is_inference=is_inference)
        # Four Type C inception modules
        for n in (128, 160, 160, 192):
            v = self.module_c(model, inputs=v, input_depth=768, n=n, is_inference=is_inference)
        # One Type D inception module
        v = self.module_d(model, inputs=v, input_depth=768, is_inference=is_inference)
        # Two Type E inception modules
        v = self.module_e(model, inputs=v, input_depth=1280, pooltype='avg', is_inference=is_inference)
        v = self.module_e(model, inputs=v, input_depth=2048, pooltype='max', is_inference=is_inference)
        # Final global pooling
        v = brew.average_pool(model, v, blob_out='pool', kernel=8, stride=1, pad=0)
        # And classifier
        return self.add_head_nodes(model, v, 2048, 'classifier', loss_scale=loss_scale) 

def inception_v4_sa(self, model, inputs, input_depth, is_inference):
        branches = [
            [('conv', 96, 3, 2, 0)],
            [('max', 3, 2, 0)]
        ]
        return self.inception_module(model, 'incept_v4_sa', inputs, input_depth,
                                     branches, is_inference) 

def inception_v4_sb(self, model, inputs, input_depth, is_inference):
        branches = [
            [('conv', 64, 1, 1, 0), ('conv', 96, 3, 1, 0)],
            [('conv', 64, 1, 1, 0), ('conv', 64, (1,7), 1, (0,3)), ('conv', 64, (7,1), 1, (3,0)),
             ('conv', 96, 3, 1, 0)]
        ]
        return self.inception_module(model, 'incept_v4_sb', inputs, input_depth,
                                     branches, is_inference) 

def inception_v4_sc(self, model, inputs, input_depth, is_inference):
        branches = [
            [('conv', 192, 3, 2, 0)],
            [('max', 3, 2, 0)]
        ]
        return self.inception_module(model, 'incept_v4_sc', inputs, input_depth,
                                     branches, is_inference)

    # Reduction functions 

def inception_v4_rb(self, model, inputs, input_depth, is_inference):
        branches = [
            [('conv', 192, 1, 1, 0), ('conv', 192, 3, 2, 0)],
            [('conv', 256, 1, 1, 0), ('conv', 256, (1,7), 1, (0,3)), ('conv', 320, (7,1), 1, (3,0)),
             ('conv', 320, 3, 2, 0)],
            [('max', 3, 2, 0)],
        ]
        return self.inception_module(model, 'incept_v4_rb', inputs, input_depth,
                                     branches, is_inference) 

def inception_v4_a(self, model, inputs, input_depth, is_inference):
        branches = [
            [('conv', 96, 1, 1, 0)],
            [('conv', 64, 1, 1, 0), ('conv', 96, 3, 1, 1)],
            [('conv', 64, 1, 1, 0), ('conv', 96, 3, 1, 1), ('conv', 96, 3, 1, 1)],
            [('avg', 3, 1, 1), ('conv', 96, 1, 1, 0)]
        ]
        return self.inception_module(model, 'incept_v4_a', inputs, input_depth,
                                     branches, is_inference) 

def inception_v4_b(self, model, inputs, input_depth, is_inference):
        branches = [
            [('conv', 384, 1, 1, 0)],
            [('conv', 192, 1, 1, 0), ('conv', 224, (1,7), 1, (0,3)), ('conv', 256, (7,1), 1, (3,0))],
            [('conv', 192, 1, 1, 0), ('conv', 192, (7,1), 1, (3,0)), ('conv', 224, (1,7), 1, (0,3)),
             ('conv', 224, (7,1), 1, (3,0)), ('conv', 256, (1,7), 1, (0,3))],
            [('avg', 3, 1, 1), ('conv', 128, 1, 1, 0)]
        ]
        return self.inception_module(model, 'incept_v4_b', inputs, input_depth,
                                     branches, is_inference) 

def inception_v4_c(self, model, inputs, input_depth, is_inference):
        branches = [
            [('conv', 256, 1, 1, 0)],
            [('conv', 384, 1, 1, 0), ('conv', 256, (1,3), 1, (0,1))],
            [('share',),             ('conv', 256, (3,1), 1, (1,0))],
            [('conv', 384, 1, 1, 0), ('conv', 448, (3,1), 1, (1,0)),
             ('conv', 512, (1,3), 1, (0,1)), ('conv', 256, (1,3), 1, (0,1))],
            [('share',), ('share',), ('share',), ('conv', 256, (3,1), 1, (1,0))],
            [('avg', 3, 1, 1), ('conv', 256, 1, 1, 0)]
        ]
        return self.inception_module(model, 'incept_v4_c', inputs, input_depth,
                                     branches, is_inference) 

def forward_pass_builder(self, model, loss_scale=1.0):
        """
            This function adds the operators, layers to the network. It should return a list
            of loss-blobs that are used for computing the loss gradient. This function is
            also passed an internally calculated loss_scale parameter that is used to scale
            your loss to normalize for the number of GPUs. Signature: function(model, loss_scale)
        """
        self.counts = defaultdict(lambda: 0)
        is_inference = self.phase == 'inference'

        v = 'data'

        # Input conv modules
        v = self.conv(model, 'conv', v, input_depth=3, num_filters=32, kernel=3, stride=2, pad=0, is_inference=is_inference)
        v = self.conv(model, 'conv', v, input_depth=32, num_filters=32, kernel=3, stride=1, pad=0, is_inference=is_inference)
        v = self.conv(model, 'conv', v, input_depth=32, num_filters=64, kernel=3, stride=1, pad=1, is_inference=is_inference)
        # Stem modules
        v = self.inception_v4_sa(model, inputs=v, input_depth=64, is_inference=is_inference)
        v = self.inception_v4_sb(model, inputs=v, input_depth=160, is_inference=is_inference)
        v = self.inception_v4_sc(model, inputs=v, input_depth=192, is_inference=is_inference)
        # Four Type A modules
        for _ in xrange(4):
            v = self.inception_v4_a(model, inputs=v, input_depth=384, is_inference=is_inference)
        # One Type A Reduction module
        v = self.inception_v4_ra(model, inputs=v, input_depth=384, k=192, l=224, m=256, n=384, is_inference=is_inference)
        # Seven Type B modules
        for _ in xrange(7):
            v = self.inception_v4_b(model, inputs=v, input_depth=1024, is_inference=is_inference)
        # One Type B Reduction module
        v = self.inception_v4_rb(model, inputs=v, input_depth=1024, is_inference=is_inference)
        # Three Type C modules
        for _ in xrange(3):
            v = self.inception_v4_c(model, inputs=v, input_depth=1536, is_inference=is_inference)
        # Final global pooling
        v = brew.average_pool(model, v, blob_out='pool', kernel=8, stride=1, pad=0)
        v = brew.dropout(model, v, 'drop', ratio=0.2, is_test=is_inference)
        # And classifier
        return self.add_head_nodes(model, v, 1536, 'classifier', loss_scale=loss_scale) 

def forward_pass_builder(self, model, loss_scale=1.0):
        """
            This function adds the operators, layers to the network. It should return a list
            of loss-blobs that are used for computing the loss gradient. This function is
            also passed an internally calculated loss_scale parameter that is used to scale
            your loss to normalize for the number of GPUs. Signature: function(model, loss_scale)
        """
        is_inference = self.phase == 'inference'

        v = 'data'

        v = brew.conv(model, v, 'conv1', 3, 96, kernel=11, stride=4)
        v = brew.relu(model, v, 'relu1')
        v = brew.lrn(model, v, 'norm1', size=5, alpha=0.0001, beta=0.75)
        v = brew.max_pool(model, v, 'pool1', kernel=3, stride=2)

        v = brew.conv(model, v, 'conv2', 96, 256, kernel=5, pad=2, group=1)
        v = brew.relu(model, v, 'relu2')
        v = brew.lrn(model, v, 'norm2', size=5, alpha=0.0001, beta=0.75)
        v = brew.max_pool(model, v, 'pool2', kernel=3, stride=2)

        v = brew.conv(model, v, 'conv3', 256, 384, kernel=3, pad=1)
        v = brew.relu(model, v, 'relu3')

        v = brew.conv(model, v, 'conv4', 384, 384, kernel=3, pad=1, group=1)
        v = brew.relu(model, v, 'relu4')

        v = brew.conv(model, v, 'conv5', 384, 256, kernel=3, pad=1, group=1)
        v = brew.relu(model, v, 'relu5')
        v = brew.max_pool(model, v, 'pool5', kernel=3, stride=2)

        v = brew.fc(model, v, 'fc6', dim_in=9216, dim_out=4096)
        v = brew.relu(model, v, 'relu6')
        v = brew.dropout(model, v, 'drop6', ratio=0.5, is_test=is_inference)

        v = brew.fc(model, v, 'fc7', dim_in=4096, dim_out=4096)
        v = brew.relu(model, v, 'relu7')
        v = brew.dropout(model, v, 'drop7', ratio=0.5, is_test=is_inference)

        return self.add_head_nodes(model, v, 4096, 'fc8', loss_scale=loss_scale) 

def Add_Original_CIFAR10_Model(model, data, num_classes, image_height, image_width, image_channels):
    # Convolutional layer 1
    conv1 = brew.conv(model, data, 'conv1', dim_in=image_channels, dim_out=32, kernel=5, stride=1, pad=2)
    h,w = update_dims(height=image_height, width=image_width, kernel=5, stride=1, pad=2)
    # Pooling layer 1
    pool1 = brew.max_pool(model, conv1, 'pool1', kernel=3, stride=2)
    h,w = update_dims(height=h, width=w, kernel=3, stride=2, pad=0)
    # ReLU layer 1
    relu1 = brew.relu(model, pool1, 'relu1')
    
    # Convolutional layer 2
    conv2 = brew.conv(model, relu1, 'conv2', dim_in=32, dim_out=32, kernel=5, stride=1, pad=2)
    h,w = update_dims(height=h, width=w, kernel=5, stride=1, pad=2)
    # ReLU layer 2
    relu2 = brew.relu(model, conv2, 'relu2')
    # Pooling layer 1
    pool2 = brew.average_pool(model, relu2, 'pool2', kernel=3, stride=2)
    h,w = update_dims(height=h, width=w, kernel=3, stride=2, pad=0)
    
    # Convolutional layer 3
    conv3 = brew.conv(model, pool2, 'conv3', dim_in=32, dim_out=64, kernel=5, stride=1, pad=2)
    h,w = update_dims(height=h, width=w, kernel=5, stride=1, pad=2)
    # ReLU layer 3
    relu3 = brew.relu(model, conv3, 'relu3')
    # Pooling layer 3
    pool3 = brew.average_pool(model, relu3, 'pool3', kernel=3, stride=2)
    h,w = update_dims(height=h, width=w, kernel=3, stride=2, pad=0)
    
    # Fully connected layers
    fc1 = brew.fc(model, pool3, 'fc1', dim_in=64*h*w, dim_out=64)
    fc2 = brew.fc(model, fc1, 'fc2', dim_in=64, dim_out=num_classes)
    
    # Softmax layer
    softmax = brew.softmax(model, fc2, 'softmax')
    return softmax


# ## Test Saved Model From Part 1
# 
# ### Construct Model for Testing
# 
# The first thing we need is a model helper object that we can attach the lmdb reader to.

# In[4]:


# Create a ModelHelper object with init_params=False 

def resnet(self, model, units, num_stages, filter_list, bottle_neck=True, bn_mom=0.9,
               is_inference=False, loss_scale=1.0):
        """Return ResNet symbol of
        Parameters
        ----------
        units : list
            Number of units in each stage
        num_stages : int
            Number of stage
        filter_list : list
            Channel size of each stage
        num_classes : int
            Ouput size of symbol
        dataset : str
            Dataset type, only cifar10 and imagenet supports
        workspace : int
            Workspace used in convolution operator''
        dtype : str
            Precision (float32 or float16)
        """
        num_unit = len(units)
        assert num_unit == num_stages
        v = 'data'
        (nchannel, _, _) = self.input_shape # (nchannel, height, width)

        v = brew.conv(model, v, 'conv0', nchannel, filter_list[0], kernel=7, pad=3,
                      stride=2, no_bias=True)
        v = brew.spatial_bn(model, v, 'bn0', filter_list[0], eps=2e-5, momentum=bn_mom,
                            is_test=is_inference)
        v = brew.relu(model, v, 'relu0')
        v = brew.max_pool(model, v, 'pool0', kernel=3, stride=2, pad=1)

        dim_in = filter_list[0]
        for i in range(num_stages):
            v = self.residual_unit(model, v, dim_in, filter_list[i+1], stride=(1 if i == 0 else 2),
                                   dim_match=False,
                                   name='stage%d_unit%d' % (i + 1, 1), bottle_neck=bottle_neck,
                                   is_inference=is_inference)
            dim_in = filter_list[i+1]
            for j in range(units[i]-1):
                v = self.residual_unit(model, v, dim_in, filter_list[i+1], 1, True,
                                       name='stage%d_unit%d' % (i + 1, j + 2),
                                       bottle_neck=bottle_neck, is_inference=is_inference)

        v = brew.average_pool(model, v, 'pool1', kernel=7, global_pool=True)
        return self.add_head_nodes(model, v, dim_in, 'classifier', loss_scale=loss_scale) 

def inception_module(self, model, name, inputs, input_depth, branches, is_inference):
        """Add parallel branches from 'branhes' into current graph"""
        self.counts[name] += 1
        name = name + str(self.counts[name]-1)
        layers_outputs = []             # Outputs of each layer in each branch
        layers_input_depths = []        # Layers input depths
        for branch_id, branch in enumerate(branches):
            v = inputs                       # Inputs for a next layer in a branch
            last_layer_output_depth = input_depth
            layers_outputs.append([])        # Outputs of layers in this branch
            layers_input_depths.append([0]*len(branch))
            for layer_id, layer in enumerate(branch):
                layers_input_depths[-1][layer_id] = last_layer_output_depth
                if layer[0] == 'conv':
                    v = self.conv(
                        model,
                        name='%s_b%d_l%d_' % (name, branch_id, layer_id),
                        inputs=v,
                        input_depth=last_layer_output_depth,
                        num_filters=layer[1],
                        kernel=layer[2],
                        stride=layer[3],
                        pad=layer[4],
                        is_inference=is_inference
                    )
                    last_layer_output_depth = layer[1]
                elif layer[0] == 'avg' or layer[0] == 'max':
                    pool_func = brew.average_pool if layer[0] == 'avg' else brew.max_pool
                    v = pool_func(
                        model, v,
                        blob_out='%s_b%d_p%d' % (name, branch_id, layer_id),
                        kernel=layer[1], stride=layer[2], pad=layer[3]
                    )
                    # last_layer_output_depth does not change here
                elif layer[0] == 'share':
                    v = layers_outputs[-2][layer_id]
                    last_layer_output_depth = layers_input_depths[-2][layer_id+1]
                else:
                    assert False, 'Unknown later type - ' + layer[0]
                layers_outputs[-1].append(v)
        # concat
        concat = brew.concat(
            model,
            [outputs[-1] for outputs in layers_outputs],
            blob_out='%s_concat' % name
        )
        return concat