Shows how to import CoreML files using F# and the CoreML Protocol Buffers

CoreMLImport.fs

module Neural.CoreMLSupport

open Data
open Neural.Layers
open CoreML.Specification
open FSharp.NativeInterop

#nowarn "9"

let rec modelDataFromCoreML (min : System.IO.Stream) : ModelData =
    //let modelFile = "/Users/fak/Downloads/my_model.mlmodel"
    //let modelFile = "/Users/fak/Downloads/Resnet50.mlmodel"
    //let modelFile = "/Users/fak/Downloads/DeepLabV3.mlmodel"
    //let modelFile = "/Users/fak/Dropbox/Projects/ImageClassifier/ImageClassifier/Resources/MobileNetV2.mlmodel"
    //let modelFile = "/Users/fak/Dropbox/Projects/Neural/Research/TwistySub.mlmodel"
    //use min = System.IO.File.OpenRead modelFile
    let m = CoreML.Specification.Model.Parser.ParseFrom min
    let layers, inputPreproc, inputArrayShape, inputImageShape, classLabels =
        match m.NeuralNetworkClassifier with
        | null ->
            match m.NeuralNetworkRegressor with
            | null ->
                match m.NeuralNetwork with
                | null -> failwithf "No neural network in the CoreML model"
                | nn -> nn.Layers, nn.Preprocessing, nn.ArrayInputShapeMapping, nn.ImageInputShapeMapping, Array.empty
            | nn -> nn.Layers, nn.Preprocessing, nn.ArrayInputShapeMapping, nn.ImageInputShapeMapping, Array.empty
        | nn ->
            let labels =
                match nn.Int64ClassLabels|>Option.ofObj, nn.StringClassLabels|>Option.ofObj with
                | _, Some c -> c.Vector |> Array.ofSeq
                | Some c, _ -> c.Vector |> Seq.map string |> Array.ofSeq
                | _ -> Array.empty
            nn.Layers, nn.Preprocessing, nn.ArrayInputShapeMapping, nn.ImageInputShapeMapping, labels
    let inputs = m.Description.Input
    let outputs = m.Description.Output
    printfn "-----------------"
    printfn "INPUTS  = %A" inputs
    printfn "OUTPUTS = %A" outputs
    printfn "PREPROC = %A" inputPreproc
    printfn "INARRAY = %A" inputArrayShape
    printfn "INIMAGE = %A" inputImageShape
    printfn "CLASSES = %A" classLabels
    //for l in layers do
    //    printfn "%O\t%s\tIN: %A\tOUT: %A" l.LayerCase l.Name l.Input l.Output

    let data = newModelData ()

    let minputsMaps =
        inputs
        |> Seq.mapi (fun i x ->
            let iname = x.Name
            let pp = inputPreproc |> Seq.tryPick (fun x -> if iname = x.FeatureName then Some x else None)
            x, { convertInputLayer i pp x with Name = (if x.Name = null then "" else x.Name) })
        |> Array.ofSeq
    let minputs = minputsMaps |> Array.map snd
    let minputIds, data = insertAll minputs data

    let moutputsMaps =
        outputs
        |> Seq.mapi (fun i x -> x, { convertOutputLayer i x classLabels with Name = (if x.Name = null then "" else x.Name) })
        |> Array.ofSeq
    let moutputs = moutputsMaps |> Array.map snd
    let moutputIds, data = insertAll moutputs data

    let mlayersMaps =
        layers
        |> Seq.mapi (fun i x -> x, { convertLayer i x with Name = (if x.Name = null then "" else x.Name) })
        |> Array.ofSeq
    let mlayers = mlayersMaps |> Array.map snd
    let mlayerIds, data = insertAll mlayers data

    let minstsMaps =
        let numCols = 10
        
        Seq.concat [ (minputsMaps |> Seq.mapi (fun i (x, y) -> x.Name, minputIds.[i], None))
                     (mlayersMaps |> Seq.mapi (fun i (x, y) -> x.Name, mlayerIds.[i], Some x))
                     (moutputsMaps |> Seq.mapi (fun i (x, y) -> x.Name, moutputIds.[i], None)) ]
        |> Seq.mapi (fun i (name, layer, nnlayer) ->
            let inst =
                {
                    Layer = layer
                    Trainable = true
                    Frame = { X = (double (i % numCols)) * 175.0; Y = (double (i / numCols)) * 200.0; Width = 100.0; Height = 100.0 }
                    DisplayValue = None
                    DisplayShape = None
                    CompileErrors = Array.empty
                }
            nnlayer, inst)
        |> Array.ofSeq
    let minsts = minstsMaps |> Array.map snd
    let minstIds, data = insertAll minsts data

    let newOutMap = System.Collections.Generic.Dictionary<string, _> ()
    for i, inp in minputsMaps |> Seq.indexed do
        let name = (fst inp).Name
        newOutMap.Add(name, (Array.empty, minstIds.[i]))
    for i, inp in mlayersMaps |> Seq.indexed do
        let outputs = (fst inp).Output
        if outputs.Count > 0 then
            let name = (fst inp).Output.[0]
            newOutMap.Add(name, ((fst inp).Input|>Array.ofSeq, minstIds.[i + minputsMaps.Length]))

    let conns = ResizeArray<Connection> ()
    for x in newOutMap do
        let thisInst = snd x.Value
        let inputNames = fst x.Value
        if inputNames.Length > 0 then
            for i, iname in inputNames |> Seq.indexed do
                let s = { Instance = snd newOutMap.[iname]; Port = "O" }
                let d = { Instance = thisInst; Port = if i > 0 then sprintf "I%d" i else "I" }
                let c = { Source = s; Destination = d }
                conns.Add(c)
    for i, outp in moutputsMaps |> Seq.indexed do
        let name = (fst outp).Name
        match newOutMap.TryGetValue name with
        | true, (_, outInst) ->
            let s = { Instance = outInst; Port = "O" }
            let d = { Instance = minstIds.[i+inputs.Count+layers.Count]; Port = "I" }
            let c = { Source = s; Destination = d }
            conns.Add(c)
        | _ -> ()

    let cids, data = insertAll (conns.ToArray()) data

    let model = { data.Root with Instances = minstIds; Connections = cids }
    let data = updateRoot model data

    let data = optimizeModel (Id data.RootId) data
    validateModelData (Id data.RootId) data
    data

and convertInputLayer (i : int) (pp : NeuralNetworkPreprocessing option) (desc : FeatureDescription) : Layer =
    let typ = convertInputFeatureType pp desc
    let kind = InputLayer { InputIndex = i; Kind = typ }
    newLayer kind

and convertOutputLayer (i : int) (desc : FeatureDescription) (classLabels : string[]) : Layer =
    let typ = convertOutputFeatureType None desc classLabels
    let kind = OutputLayer { OutputIndex = i; Kind = typ }
    newLayer kind

and convertImageType (pp : NeuralNetworkPreprocessing option) (img : ImageFeatureType) =
    let cs =
        match img.ColorSpace with
        | ImageFeatureType.Types.ColorSpace.Bgr -> BgrColorSpace
        | ImageFeatureType.Types.ColorSpace.Grayscale -> GrayColorSpace
        | _ -> RgbColorSpace
    let scale, rb, gb, bb =
        match pp with
        | Some pp when pp.PreprocessorCase = NeuralNetworkPreprocessing.PreprocessorOneofCase.Scaler ->
            pp.Scaler.ChannelScale, pp.Scaler.RedBias, pp.Scaler.GreenBias, pp.Scaler.BlueBias
        | _ -> (1.0f / 127.5f), -1.0f, -1.0f, -1.0f
    let c =
        { 
            ColorSpace = cs
            ChannelScale = scale
            RedBias = rb
            GreenBias = gb
            BlueBias = bb
        }
    (img.Width |> Expr.ofInt64), (img.Height |> Expr.ofInt64), c

and convertInputFeatureType (pp : NeuralNetworkPreprocessing option) (desc : FeatureDescription) : InputKind<Expr> =
    match desc.Type.TypeCase with
    | FeatureType.TypeOneofCase.ImageType ->
        let w, h, c = convertImageType pp desc.Type.ImageType
        ImageInput {Width=w;Height=h;Channels=c}
    | _ ->
        ArrayInput [| Expr.oneInt |]

and convertOutputFeatureType (pp : NeuralNetworkPreprocessing option) (desc : FeatureDescription) (classLabels : string[]) : OutputKind =
    match desc.Type.TypeCase, classLabels with
    | FeatureType.TypeOneofCase.ImageType, _ ->
        let _, _, c = convertImageType pp desc.Type.ImageType
        ImageOutput c
    | _, [||] -> ArrayOutput
    | _, _ -> OneHotOutput classLabels

and convertLayer (i : int) (layer : NeuralNetworkLayer) : Layer =
    let kind, weights =
        match layer.LayerCase with
        | NeuralNetworkLayer.LayerOneofCase.Activation ->
            let a =
                match layer.Activation.NonlinearityTypeCase with
                | ActivationParams.NonlinearityTypeOneofCase.None -> IdentityActivation
                | ActivationParams.NonlinearityTypeOneofCase.ReLU -> ReLUActivation
                | ActivationParams.NonlinearityTypeOneofCase.LeakyReLU -> LeakyReLUActivation (layer.Activation.LeakyReLU.Alpha |> Expr.ofFloat)
                | ActivationParams.NonlinearityTypeOneofCase.ThresholdedReLU -> ThresholdActivation (layer.Activation.ThresholdedReLU.Alpha |> Expr.ofFloat, Expr.zeroFloat)
                | ActivationParams.NonlinearityTypeOneofCase.Tanh -> TanhActivation
                | ActivationParams.NonlinearityTypeOneofCase.Linear ->
                    match layer.Activation.Linear.Alpha, layer.Activation.Linear.Beta with
                    | 1.0f, 0.0f -> IdentityActivation
                    | a, b -> LinearActivation (a |> Expr.ofFloat, b |> Expr.ofFloat)
                | x ->
                    printfn "UNKNOWN ACTIVATION %A" x
                    LinearActivation (Expr.oneFloat, Expr.zeroFloat)
            ActivationLayer a, Map.empty
        | NeuralNetworkLayer.LayerOneofCase.AddBroadcastable
        | NeuralNetworkLayer.LayerOneofCase.Add ->
            BinopLayer "+", Map.empty
        | NeuralNetworkLayer.LayerOneofCase.Batchnorm ->
            let weights =
                seq {
                    if layer.Batchnorm.Beta <> null then
                        "Beta", weightsToTensor layer.Batchnorm.Beta
                    if layer.Batchnorm.Gamma <> null then
                        "Gamma", weightsToTensor layer.Batchnorm.Gamma
                    if layer.Batchnorm.Mean <> null then
                        "Mean", weightsToTensor layer.Batchnorm.Mean
                    if layer.Batchnorm.Variance <> null then
                        "Variance", weightsToTensor layer.Batchnorm.Variance
                }
                |> Map.ofSeq
            let layerc =
                { NormalizationConfig<Expr>.DefaultBatchNorm with
                    Epsilon = layer.Batchnorm.Epsilon |> Expr.ofFloat
                }
            NormalizationLayer layerc, weights
        | NeuralNetworkLayer.LayerOneofCase.Concat ->
            let c = layer.Concat
            if c.SequenceConcat then
                failwith "Sequence cat not supported"
            ConcatenationLayer { ConcatenationAxis = Expr.oneInt }, Map.empty
        | NeuralNetworkLayer.LayerOneofCase.Convolution ->
            let conv = layer.Convolution
            if conv.IsDeconvolution then
                failwith "Deconvolution not supported"
            let weights =
                match conv.Weights with
                | null -> Map.empty
                | weights ->
                    let tweights = weightsToTensor weights
                    match conv.HasBias, layer.Convolution.Bias with
                    | _, null
                    | false, _ -> seq{"Weights", tweights}|>Map.ofSeq
                    | true, bias ->
                        let tbias = weightsToTensor bias
                        seq{"Weights", tweights;"Biases", tbias}|>Map.ofSeq
            ConvolutionLayer
                {
                    Filters = int conv.OutputChannels |> Expr.ofInt
                    KernelSizes = conv.KernelSize |> toExprArray
                    Strides = conv.Stride |> toExprArray
                    Dilation = conv.DilationFactor |> toExprArray
                    UseBias = conv.HasBias
                    Groups = int conv.NGroups |> Expr.ofInt
                    Padding =
                        match conv.ConvolutionPaddingTypeCase with
                        | ConvolutionLayerParams.ConvolutionPaddingTypeOneofCase.Valid -> Neural.Layers.ValidPadding
                        | ConvolutionLayerParams.ConvolutionPaddingTypeOneofCase.Same
                        | _ -> Neural.Layers.SamePadding
                    Activation = None
                    Normalization = None
                    Bias = ZerosInitializer
                    Kernel = WeightsInitializer.Uniform
                }, weights
        | NeuralNetworkLayer.LayerOneofCase.Flatten ->
            FlattenLayer { FlattenDimensions = Expr.oneInt }, Map.empty
        | NeuralNetworkLayer.LayerOneofCase.FlattenTo2D ->
            FlattenLayer { FlattenDimensions = Expr.twoInt }, Map.empty
        | NeuralNetworkLayer.LayerOneofCase.GreaterThan ->
            ActivationLayer (GreaterThan (layer.GreaterThan.Alpha |> Expr.ofFloat)), Map.empty
        | NeuralNetworkLayer.LayerOneofCase.InnerProduct ->
            DenseLayer { DenseConfig<Expr>.Default with Units = layer.InnerProduct.OutputChannels |> Expr.ofUInt64 }, Map.empty
        | NeuralNetworkLayer.LayerOneofCase.LoadConstant ->
            let c = layer.LoadConstant
            let shape = c.Shape
            let weights =
                match c.Data with
                | null -> Map.empty
                | v -> seq{"Value", weightsToTensor v}|>Map.ofSeq
            let c =
                {
                    VariableShape = shape |> toExprArray
                    Initializer = WeightsInitializer.Uniform
                }
            VariableLayer c, weights
        | NeuralNetworkLayer.LayerOneofCase.Multiply ->
            BinopLayer "*", Map.empty
        | NeuralNetworkLayer.LayerOneofCase.Padding ->
            match layer.Padding.PaddingTypeCase with
            | PaddingLayerParams.PaddingTypeOneofCase.Constant when layer.Padding.PaddingAmounts.BorderAmounts_.Count = 2 ->
                let h = layer.Padding.PaddingAmounts.BorderAmounts_.[0]
                let v = layer.Padding.PaddingAmounts.BorderAmounts_.[1]
                PaddingLayer (sprintf "%dx%d = %g" h.EndEdgeSize v.EndEdgeSize layer.Padding.Constant.Value), Map.empty
            | _ ->
                PaddingLayer (sprintf "%O" layer.Padding.PaddingTypeCase), Map.empty
        | NeuralNetworkLayer.LayerOneofCase.Permute ->
            let p = layer.Permute
            TransposeLayer { TransposeAxes = p.Axis |> toExprArray; IsPermutation = true }, Map.empty
        | NeuralNetworkLayer.LayerOneofCase.Pooling ->
            let pooling = layer.Pooling
            let kind =
                match pooling.Type with
                | PoolingLayerParams.Types.PoolingType.Average -> AveragePooling (not pooling.AvgPoolExcludePadding)
                | PoolingLayerParams.Types.PoolingType.L2 -> L2NormPooling
                | PoolingLayerParams.Types.PoolingType.Max | _ -> MaxPooling
            let c =
                {
                    KernelSizes = pooling.KernelSize |> toExprArray
                    Padding =
                        match pooling.PoolingPaddingTypeCase with
                        | PoolingLayerParams.PoolingPaddingTypeOneofCase.Valid -> Neural.Layers.ValidPadding
                        | PoolingLayerParams.PoolingPaddingTypeOneofCase.Same
                        | _ -> Neural.Layers.SamePadding
                    Strides = pooling.Stride |> toExprArray
                    Global = pooling.GlobalPooling
                    Kind = kind
                }
            PoolingLayer c, Map.empty
        | NeuralNetworkLayer.LayerOneofCase.Reshape ->
            let r = layer.Reshape
            let mode =
                match r.Mode with
                | ReshapeLayerParams.Types.ReshapeOrder.ChannelFirst -> ChannelsFirst
                | _ -> ChannelsLast
            let shape = r.TargetShape |> toInt64ExprArray
            ReshapeLayer {Reshape = shape; Mode=mode}, Map.empty
        | NeuralNetworkLayer.LayerOneofCase.ResizeBilinear ->
            let ysize = layer.ResizeBilinear.TargetSize.[0]
            let xsize = layer.ResizeBilinear.TargetSize.[1]
            let scale = TargetSize (xsize |> Expr.ofUInt64, ysize |> Expr.ofUInt64)
            let align =
                match layer.ResizeBilinear.Mode.SamplingMethod with
                | SamplingMode.Types.Method.StrictAlignEndpointsMode -> true
                | SamplingMode.Types.Method.AlignEndpointsMode -> true
                | _ -> false
            UpsampleLayer { UpScale=scale; Sampling=LinearSampling; AlignCorners=align }, Map.empty
        | NeuralNetworkLayer.LayerOneofCase.Reduce ->
            let r = layer.Reduce
            let kind =
                match r.Mode with
                | ReduceLayerParams.Types.ReduceOperation.Argmax -> ReductionKind.ArgMaxReduction
                //| ReduceLayerParams.Types.ReduceOperation.Argmin -> ReductionKind.ArgMinReduction
                | ReduceLayerParams.Types.ReduceOperation.Max -> ReductionKind.MaxReduction
                | ReduceLayerParams.Types.ReduceOperation.Avg -> ReductionKind.MeanReduction
                | ReduceLayerParams.Types.ReduceOperation.Min -> ReductionKind.MinReduction
                | ReduceLayerParams.Types.ReduceOperation.Sum -> ReductionKind.SumReduction
                | ReduceLayerParams.Types.ReduceOperation.L1 -> ReductionKind.L1NormReduction
                | _ -> failwithf "Reduction type not supported: %O" r.Mode
            let axes =
                match r.Axis with
                | ReduceLayerParams.Types.ReduceAxis.C -> [| 1 |]
                | ReduceLayerParams.Types.ReduceAxis.Chw -> [| 1; 2; 3 |]
                | ReduceLayerParams.Types.ReduceAxis.H -> [| 2 |]
                | ReduceLayerParams.Types.ReduceAxis.Hw -> [| 2; 3 |]
                | ReduceLayerParams.Types.ReduceAxis.W -> [| 3 |]
                | _ -> failwithf "Reduction axis not supported: %O" r.Axis
            ReductionLayer { ReductionAxes = axes |> toIntExprArray; Kind = kind }, Map.empty
        | NeuralNetworkLayer.LayerOneofCase.Softmax ->
            ActivationLayer (SoftmaxActivation Expr.oneInt), Map.empty
        | NeuralNetworkLayer.LayerOneofCase.Transpose ->
            TransposeLayer { TransposeAxes = layer.Transpose.Axes |> toExprArray; IsPermutation = false }, Map.empty
        | NeuralNetworkLayer.LayerOneofCase.Unary ->
            //printfn "UNKNOWN UNARY %O" layer.Unary.Type
            match layer.Unary.Type with
            | UnaryFunctionLayerParams.Types.Operation.Threshold ->
                UnopLayer (ThresholdUnop (layer.Unary.Alpha |> Expr.ofFloat)), Map.empty
            | _ -> UnopLayer (UnknownUnop (sprintf "%O a=%g" layer.Unary.Type layer.Unary.Alpha)), Map.empty
        | _ ->
            failwithf "Unsupported CoreML Layer: %O" layer.LayerCase
            //DenseLayer DenseConfig<Expr>.Default, Map.empty
    { newLayer kind with Weights = weights }

and private toShape x : Shape =
    x |> Seq.map int |> Shape.ofSeq

and private toIntArray (x : uint64 seq) : int[] =
    x |> Seq.map int |> Array.ofSeq

and private toExprArray (x : uint64 seq) : Expr[] =
    x |> Seq.map Expr.ofUInt64 |> Array.ofSeq

and private toIntExprArray (x : int seq) : Expr[] =
    x |> Seq.map Expr.ofInt |> Array.ofSeq

and private toInt64ExprArray (x : int64 seq) : Expr[] =
    x |> Seq.map Expr.ofInt64 |> Array.ofSeq

and weightsToTensor (w : WeightParams) : Neural.Tensor =
    match w.FloatValue with
    | f when f <> null && f.Count > 0 ->
        let n = f.Count
        let shape = [|1;n;1;1|]
        let t = Neural.Tensor.Fill (0.0f, shape)
        let fp = t.Floats 
        for i in 0 .. (n - 1) do
            NativePtr.set fp i f.[i]
        t
    | _ -> Neural.Tensor.Fill (0.0f, [|1;1;1;1|])