TorchSharpLeak2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "dotnet_interactive": {
     "language": "fsharp"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "Installed package libtorch-cpu version 1.8.0.7"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "Installed package TorchSharp version 0.91.52518"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "TorchSharp: LoadNativeBackend: Native backend not found in application loading TorchSharp directly from packages directory.\n"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "TorchSharp: LoadNativeBackend: Trying dynamic load for .NET/F# Interactive by consolidating native libtorch-cpu-* binaries to /home/peter/.nuget/packages/torchsharp/0.91.52518/lib/netcoreapp3.1/cpu...\n"
     },
     "output_type": "unknown"
    }
   ],
   "source": [
    "#r \"nuget: libtorch-cpu, 1.8.0.7\"\n",
    "#r \"nuget: TorchSharp, 0.91.52518\"\n",
    "\n",
    "open TorchSharp\n",
    "open TorchSharp.Tensor\n",
    "open TorchSharp.NN\n",
    "let device = Torch.InitializeDevice Device.CPU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### PCA + Deconvolution\n",
    "\n",
    "We render 1000s of traces with 1D random `xs` and calculated `ys` and wish to infer parameters of transform.\n",
    "\n",
    "The transform is done by:\n",
    "- choosing a random mixture of 4 component functions (e.g. sin(x), cos(x))\n",
    "- adding the mixture into a kernel\n",
    "- use the kernel to convolve random vector of `xs` into `ys`\n",
    "\n",
    "The task is to reconstruct the:\n",
    "- shape of 4 components inside kernels\n",
    "- mixture weights for 4 components for each sample (embedding lookups)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "dotnet_interactive": {
     "language": "fsharp"
    }
   },
   "outputs": [],
   "source": [
    "// helper\n",
    "type TorchTensor with\n",
    "    member t.toArray() =\n",
    "        match t.shape with\n",
    "        | [|n|] -> Array.init (int n) (fun i -> t.[int64 i].ToSingle())\n",
    "        | _ -> failwithf \"requires 1-dimensional tensor, got %A\" t.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "dotnet_interactive": {
     "language": "fsharp"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": "<table><thead><tr><th>Item1</th><th>Item2</th></tr></thead><tbody><tr><td><div class=\"dni-plaintext\">[ 984, 252, 659, 0, 879, 585, 136, 158, 556, 803, 477, 161, 876, 853, 37, 339, 312, 359, 641, 73 ... (12 more) ]</div></td><td><div class=\"dni-plaintext\">[ 8, 7, 8, 8, 8, 7, 7, 7, 7, 7, 6, 7, 7, 6, 6, 7, 7 ]</div></td></tr></tbody></table>"
     },
     "output_type": "unknown"
    }
   ],
   "source": [
    "let xsLen = 32\n",
    "let convLen = 16\n",
    "let components = [| // component functions that we will later try to reconstruct (deconvolve)\n",
    "    fun x -> (Math.Sin (x/float convLen*3.14) * 0.5 + 0.5)\n",
    "    fun x -> (Math.Cos (x/float convLen*3.14) * 0.5 + 0.5)\n",
    "    fun x -> (-Math.Sin (x/float convLen*3.14) * 0.5 + 0.5)\n",
    "    fun x -> (-Math.Cos (x/float convLen*3.14) * 0.5 + 0.5)\n",
    "|]\n",
    "let nComponents = components.Length\n",
    "\n",
    "let random = Random()\n",
    "module SampleGenerator =\n",
    "    let dirichlet n =\n",
    "        let ps = Array.init n (fun _ -> random.NextDouble()**2.0)\n",
    "        let sum = ps |> Array.sum\n",
    "        ps |> Array.map (fun p -> p / sum)\n",
    "    let randomComponentMixtureKernel len =\n",
    "        let mixture = dirichlet components.Length\n",
    "        let compMix = Array.zip components mixture\n",
    "        Array.init len (fun i -> compMix |> Array.sumBy (fun (compFn,weight) -> (compFn (float i) * weight)))\n",
    "    let conv (kernel: float[]) (xs: float[]) =\n",
    "        Array.init (xs.Length-kernel.Length+1) (fun i ->\n",
    "            let mutable sum = 0.0\n",
    "            for j in 0 .. kernel.Length-1 do sum <- sum + xs.[i+j] * kernel.[j]\n",
    "            sum)\n",
    "    let flipCoin (rate: float) (trials: float) = Math.Round(trials * rate)\n",
    "    let renderSample xsLen convLen =\n",
    "        let kernel = randomComponentMixtureKernel convLen\n",
    "        let xs = Array.init xsLen (fun _ -> random.Next 1000 |> float)\n",
    "        let rate = random.NextDouble() * 0.01\n",
    "        let ys = xs |> conv kernel |> Array.map (flipCoin rate)\n",
    "        xs, ys\n",
    "\n",
    "SampleGenerator.renderSample xsLen convLen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "dotnet_interactive": {
     "language": "fsharp"
    }
   },
   "outputs": [],
   "source": [
    "let nTraces = 10000\n",
    "\n",
    "let renderDataset xsLen convLen nTraces =\n",
    "    let tensor (xs: float[]) = Float32Tensor.from(Array.map float32 xs, false)\n",
    "    let xs, ys =\n",
    "        Array.init nTraces (fun _ -> \n",
    "            let xs, ys = SampleGenerator.renderSample xsLen convLen\n",
    "            tensor xs, tensor ys)\n",
    "        |> Array.unzip\n",
    "    let mutable i0=0\n",
    "    fun batchSize ->\n",
    "        let xs, ys, indices =\n",
    "            Array.init batchSize (fun i -> \n",
    "                i0 <- (i0+1) % nTraces\n",
    "                xs.[i0], ys.[i0], i0)\n",
    "            |> Array.unzip3\n",
    "        xs.stack 0L, ys.stack 0L, Int32Tensor.from(indices, false)\n",
    "\n",
    "let generateBatch = renderDataset xsLen convLen nTraces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "dotnet_interactive": {
     "language": "fsharp"
    }
   },
   "outputs": [],
   "source": [
    "let inline fscalar x = TorchScalar.op_Implicit (float32 x)\n",
    "let epsilon = fscalar 10e-12\n",
    "\n",
    "type Model(device, nTraces, convLen, nComponents) =\n",
    "    inherit CustomModule(\"deconv\")\n",
    "    let nTraces, nComponents, convLen = int64 nTraces, int64 nComponents, int64 convLen\n",
    "    let logKernel = Float32Tensor.rand([|nComponents; 1L; convLen|], device, true)\n",
    "    let logScale = Float32Tensor.from([|-7.0f|], true)\n",
    "    let logEmbeddings = TorchSharp.NN.Modules.Embedding( nTraces, nComponents )\n",
    "\n",
    "    member this.parameters = [| logKernel; logScale; logEmbeddings.Weight |]\n",
    "\n",
    "    override _.forward (x:TorchTensor) = failwithf \"wrong method\"\n",
    "\n",
    "    member n.forward (xs:TorchTensor, indices:TorchTensor) =\n",
    "        let factors = (logEmbeddings.forward indices + logScale).exp().unsqueeze(2L)\n",
    "        let kernel = logKernel.exp()\n",
    "        let ins = xs.unsqueeze(1L).expand([|-1L; nComponents; -1L;|])\n",
    "        //printfn \"ins=%A kernel=%A\" ins.shape kernel.shape\n",
    "        let compOuts = ins.conv1d(kernel, groups=nComponents)\n",
    "        //printfn \"conv=%A, factors=%A, kernel=%A\" compOuts.shape factors.shape kernel.shape\n",
    "        let outs = (compOuts * factors).sum([|1L|], keepDimension=false) + epsilon //* globalFactors.exp()\n",
    "        outs\n",
    "\n",
    "    member _.Kernel with get () = logKernel.exp()\n",
    "    member _.Scale with get () = logScale\n",
    "    member n.modelLoss() =\n",
    "        (logKernel.exp().sum([|2L|], keepDimension=true) - fscalar (convLen/2L)).abs().mean()\n",
    "        //+ (logEmbeddings.Weight.exp().sum([|1L|], keepDimension=true) - fscalar 1.0).mean()\n",
    "\n",
    "let net = new Model(device, nTraces, convLen, nComponents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "dotnet_interactive": {
     "language": "fsharp"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "[|3.0f; 3.0f; 3.0f; 3.0f; 3.0f; 3.0f; 3.0f; 3.0f; 3.0f; 3.0f; 3.0f; 3.0f; 3.0f;\n  3.0f; 3.0f; 3.0f; 3.0f|] -> [|129.3888245f; 130.4532471f; 124.7322311f; 114.5003128f; 108.2777328f;\n  108.2443237f; 105.8801346f; 91.25262451f; 94.07717896f; 101.5460815f;\n  102.5972824f; 111.5687332f; 110.2858582f; 115.1255493f; 120.9230194f;\n  121.8119354f; 108.3330994f|]"
     },
     "output_type": "unknown"
    }
   ],
   "source": [
    "let xs, ys, indices = generateBatch 5\n",
    "let ys' = net.forward(xs,indices)\n",
    "//sprintf \"shapes: xs=%A ys=%A indices=%A -> ys'=%A\" xs.shape ys.shape indices.shape ys'.shape\n",
    "sprintf \"%A -> %A\" (ys.[0L].toArray()) (ys'.[0L].toArray())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "dotnet_interactive": {
     "language": "fsharp"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": "<div class=\"dni-plaintext\">201.26754760742188</div>"
     },
     "output_type": "unknown"
    }
   ],
   "source": [
    "let xlogy(x:TorchTensor, y:TorchTensor) = (x.clamp_min epsilon) * y.log()\n",
    "\n",
    "let poissonLoss (k:TorchTensor) (mu:TorchTensor) =\n",
    "    let logPmf = xlogy(k,mu) - (k+fscalar 1.0).lgamma() - mu\n",
    "    -logPmf\n",
    "\n",
    "let criterion ys ys' = (poissonLoss ys ys').clamp_max(fscalar 10000.0).mean()\n",
    "//let criterion (ys:TorchTensor) (ys':TorchTensor) = let d = ys - ys' in (d*d).mean()\n",
    "\n",
    "let xs, ys, indices = generateBatch 1\n",
    "let ys' = net.forward(xs,indices)\n",
    "//printfn \"xs.shape=%A ys'.shape=%A, result=%A\" xs.shape ys'.shape (ys'.[0L].toArray())\n",
    "(criterion ys ys').ToDouble()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "dotnet_interactive": {
     "language": "fsharp"
    }
   },
   "outputs": [],
   "source": [
    "let optimizer = TorchSharp.NN.Optimizer.Adam(net.parameters, 0.02)\n",
    "//let optimizer = TorchSharp.NN.Optimizer.SGD(net.parameters, 0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "dotnet_interactive": {
     "language": "fsharp"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "step      0: loss=3.0486 loss0=3.0486 scale=-49.2762"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step  10000: loss=2.3708 loss0=2.3075 scale=-49.6194"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step  20000: loss=2.3694 loss0=2.4371 scale=-49.6499"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step  30000: loss=2.3691 loss0=2.4545 scale=-49.9243"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step  40000: loss=2.3697 loss0=2.3071 scale=-50.1237"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step  50000: loss=2.3693 loss0=2.4524 scale=-50.1515"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step  60000: loss=2.3696 loss0=2.3967 scale=-50.3999"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step  70000: loss=2.3691 loss0=2.3694 scale=-50.6440"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step  80000: loss=2.3698 loss0=2.5228 scale=-50.6729"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step  90000: loss=2.3693 loss0=2.3190 scale=-50.8671"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step 100000: loss=2.3696 loss0=2.3835 scale=-51.1481"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step 110000: loss=2.3695 loss0=2.4214 scale=-51.2014"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step 120000: loss=2.3700 loss0=2.3012 scale=-51.3203"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step 130000: loss=2.3699 loss0=2.4819 scale=-51.6288"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step 140000: loss=2.3700 loss0=2.3195 scale=-51.7434"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step 150000: loss=2.3696 loss0=2.4119 scale=-51.7996"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step 160000: loss=2.3691 loss0=2.4656 scale=-52.0932"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step 170000: loss=2.3694 loss0=2.2967 scale=-52.2724"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step 180000: loss=2.3698 loss0=2.4154 scale=-52.3066"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step 190000: loss=2.3702 loss0=2.5066 scale=-52.5864"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "step 200000: loss=2.3697 loss0=2.2861 scale=-52.7874"
     },
     "output_type": "unknown"
    }
   ],
   "source": [
    "let batchSize = 384\n",
    "let mutable cumLoss = 0.0\n",
    "let mutable nItems = 0\n",
    "for i in 0..200000 do\n",
    "    net.ZeroGrad()\n",
    "    optimizer.zero_grad()\n",
    "    let xs, ys, indices = generateBatch batchSize\n",
    "    let ys' = net.forward(xs,indices)\n",
    "    let loss = criterion ys ys' + net.modelLoss()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    cumLoss <- cumLoss + loss.ToDouble()\n",
    "    nItems <- nItems + 1\n",
    "    if i%10000 = 0 then\n",
    "        System.Console.Write $\"step %6d{i}: loss=%.4f{cumLoss / float nItems} loss0=%.4f{loss.ToDouble()} scale=%.4f{net.Scale.ToDouble()}\"\n",
    "        cumLoss <- 0.0\n",
    "        nItems <- 0\n",
    "    if i%1000 = 0 then\n",
    "        GC.Collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "dotnet_interactive": {
     "language": "fsharp"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": "<div><strong>Restore sources</strong><ul><li><span>https://pkgs.dev.azure.com/dnceng/public/_packaging/dotnet5/nuget/v3/index.json</span></li><li><span>https://pkgs.dev.azure.com/dnceng/public/_packaging/dotnet-tools/nuget/v3/index.json</span></li></ul></div>"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "Installed package Plotly.NET version 2.0.0-beta9"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/plain": "Installed package Plotly.NET.Interactive version 2.0.0-beta9"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/markdown": "Loading extensions from `Plotly.NET.Interactive.dll`"
     },
     "output_type": "unknown"
    },
    {
     "data": {
      "text/markdown": "Added Kernel Extension including formatters for GenericChart"
     },
     "output_type": "unknown"
    }
   ],
   "source": [
    "#i \"nuget:https://pkgs.dev.azure.com/dnceng/public/_packaging/dotnet5/nuget/v3/index.json\"\n",
    "#i \"nuget:https://pkgs.dev.azure.com/dnceng/public/_packaging/dotnet-tools/nuget/v3/index.json\"\n",
    "\n",
    "#r \"nuget: Plotly.NET, 2.0.0-beta9\"\n",
    "#r \"nuget: Plotly.NET.Interactive, 2.0.0-beta9\"\n",
    "open Plotly.NET\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "dotnet_interactive": {
     "language": "fsharp"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": "<div id=\"2065c4f8-b9ca-40b2-bc4e-8c8d79e96f76\" style=\"width: 1000px; height: 600px;\"><!-- Plotly chart will be drawn inside this DIV --></div>\n<script type=\"text/javascript\">\n\n            var renderPlotly_2065c4f8b9ca40b2bc4e8c8d79e96f76 = function() {\n            var fsharpPlotlyRequire = requirejs.config({context:'fsharp-plotly',paths:{plotly:'https://cdn.plot.ly/plotly-latest.min'}}) || require;\n            fsharpPlotlyRequire(['plotly'], function(Plotly) {\n\n            var data = [{\"type\":\"scatter\",\"x\":[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16],\"y\":[0.6199423,0.69567466,0.73556954,0.77553624,0.78286314,0.7736496,0.73156565,0.68186796,0.6069255,0.5234543,0.4204886,0.3082383,0.20308305,0.09811954,0.0101151075,2.1703383E-06],\"mode\":\"lines\",\"line\":{},\"marker\":{}},{\"type\":\"scatter\",\"x\":[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16],\"y\":[0.43875274,0.51339185,0.57557374,0.62894684,0.6673624,0.6811963,0.68980813,0.675424,0.6438463,0.5915692,0.5199733,0.4486551,0.3623314,0.27444956,0.18122855,0.09106393],\"mode\":\"lines\",\"line\":{},\"marker\":{}},{\"type\":\"scatter\",\"x\":[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16],\"y\":[1.4926878,1.2807386,1.0595913,0.84294283,0.6398637,0.46763194,0.30384222,0.19043581,0.106659845,0.06560287,0.069493234,0.090518706,0.16530739,0.26458606,0.40483522,0.5390727],\"mode\":\"lines\",\"line\":{},\"marker\":{}},{\"type\":\"scatter\",\"x\":[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16],\"y\":[0.000317639,0.008022388,0.054756247,0.099874645,0.16290599,0.23542237,0.32557592,0.41496727,0.5101782,0.60321206,0.69829994,0.7906289,0.8653516,0.9293591,0.97047955,0.97861814],\"mode\":\"lines\",\"line\":{},\"marker\":{}}];\n            var layout = {\"width\":1000.0,\"height\":600.0};\n            var config = {};\n            Plotly.newPlot('2065c4f8-b9ca-40b2-bc4e-8c8d79e96f76', data, layout, config);\n});\n            };\n            if ((typeof(requirejs) !==  typeof(Function)) || (typeof(requirejs.config) !== typeof(Function))) {\n                var script = document.createElement(\"script\");\n                script.setAttribute(\"src\", \"https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.6/require.min.js\");\n                script.onload = function(){\n                    renderPlotly_2065c4f8b9ca40b2bc4e8c8d79e96f76();\n                };\n                document.getElementsByTagName(\"head\")[0].appendChild(script);\n            }\n            else {\n                renderPlotly_2065c4f8b9ca40b2bc4e8c8d79e96f76();\n            }\n</script>\n"
     },
     "output_type": "unknown"
    }
   ],
   "source": [
    "let traces = [\n",
    "    for i in 0L..3L ->\n",
    "        let ys = net.Kernel.[i].[0L].toArray()\n",
    "        let xs = [|0..ys.Length|]\n",
    "        Chart.Line(xs, ys)\n",
    "]\n",
    "traces\n",
    "|> Chart.Combine\n",
    "|> Chart.withSize(1000.,600.)\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {},
  "orig_nbformat": 3
 },
 "nbformat": 4,
 "nbformat_minor": 2
}