def process_train_proto()

in mmdnn/conversion/caffe/graph.py [0:0]

• 41 lines of code
• 24 McCabe index (conditional complexity)

    def process_train_proto(self):
        layers = self.model.layer or self.model.layers
        delete_layer = set()
        split_op_map = dict()
        loss_layers = [layer for layer in layers if NodeKind.map_raw_kind(layer.type) in (NodeKind.SoftmaxWithLoss, NodeKind.SigmoidCrossEntropyLoss)]
        a = [layers.remove(layer) for layer in layers[:] if layer in loss_layers[:-1] or NodeKind.map_raw_kind(layer.type) in LAYER_IN_TRAIN_PROTO]
        for layer in layers[:]:
            if 'label' in layer.bottom:
                if NodeKind.map_raw_kind(layer.type) in (NodeKind.SoftmaxWithLoss, NodeKind.SigmoidCrossEntropyLoss):
                    continue
                elif NodeKind.map_raw_kind(layer.type) == NodeKind.Split:
                    for item in layer.top:
                        delete_layer.add(item)
                layers.remove(layer)
            elif NodeKind.map_raw_kind(layer.type) == NodeKind.Split:
                for item in layer.top:
                    split_op_map[item] = layer.bottom[0]
                layers.remove(layer)

        for layer in layers[:]:
            for item in delete_layer:
                if item in layer.bottom:
                    layers.remove(layer)
                    break
            for key, value in split_op_map.items():
                if key in layer.bottom:
                    layer.bottom.remove(key)
                    layer.bottom.append(value)
        self.model.input.append('data')
        self.model.input_dim.extend(self.input_shape)
        last_layer = layers[-1]
        kind = NodeKind.map_raw_kind(last_layer.type)
        if kind in (NodeKind.SoftmaxWithLoss, NodeKind.SigmoidCrossEntropyLoss):
            pred = layers.add()
            pred.name = 'prob'
            pred.top.append('prob')
            pred.bottom.append(last_layer.bottom[0])
            if kind == NodeKind.SoftmaxWithLoss:
                pred.type = NodeKind.Softmax if self.model.layer else 20 # competiable with old version caffe proto
            elif kind == NodeKind.SigmoidCrossEntropyLoss:
                pred.type = NodeKind.Sigmoid if self.model.layer else 19
        layers.remove(last_layer)