# Module kmath-noa A Bayesian computation library over [NOA](https://github.com/grinisrit/noa.git) together with relevant functionality from [LibTorch](https://pytorch.org/cppdocs). Our aim is to cover a wide set of applications from deep learning to particle physics simulations. In fact, we support any differentiable program written on top of `AutoGrad` & `ATen`. ## Installation from source Currently, we support only the [GNU](https://gcc.gnu.org/) toolchain for the native artifacts. For `GPU` kernels, we require a compatible [CUDA](https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html) installation. If you are on Windows, we recommend setting up everything on [WSL](https://docs.nvidia.com/cuda/wsl-user-guide/index.html). To install the library, you can simply publish to the local Maven repository: ``` ./gradlew -q :kmath-noa:publishToMavenLocal ``` This will fetch and build the `JNI` wrapper `jnoa`. In your own application add the local dependency: ```kotlin repositories { mavenCentral() mavenLocal() } dependencies { implementation("space.kscience:kmath-noa:0.3.0-dev-14") } ``` To load the native library you will need to add to the VM options: ``` -Djava.library.path=${HOME}/.konan/third-party/noa-v0.0.1/cpp-build/kmath ``` ## Usage The library spans several use cases. ###Linear Algebra We implement the tensor algebra interfaces from [kmath-tensors](../kmath-tensors): ```kotlin NoaFloat { val tensor = randNormal( shape = intArrayOf(7, 5, 3), device = Device.CPU) // or Device.CUDA(0) for GPU // Compute SVD val (tensorU, tensorS, tensorV) = tensor.svd() // Reconstruct tensor val tensorReg = tensorU dot (diagonalEmbedding(tensorS) dot tensorV.transpose(-2, -1)) } ``` ### Automatic Differentiation The [AutoGrad](https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html) engine is exposed: ```kotlin NoaFloat { // Create a quadratic function val dim = 3 val tensorX = randNormal(shape = intArrayOf(dim)) val randFeatures = randNormal(shape = intArrayOf(dim, dim)) val tensorSigma = randFeatures + randFeatures.transpose(0, 1) val tensorMu = randNormal(shape = intArrayOf(dim)) // Create a differentiable expression val expressionAtX = withGradAt(tensorX) { x -> 0.5f * (x dot (tensorSigma dot x)) + (tensorMu dot x) + 25.9f } // Evaluate the gradient at tensorX // retaining the graph for the hessian computation val gradientAtX = expressionAtX.autoGradient(tensorX, retainGraph = true) // Compute the hessian at tensorX val hessianAtX = expressionAtX.autoHessian(tensorX) } ``` ### Deep Learning You can train any [TorchScript](https://pytorch.org/docs/stable/jit.html) model. For example, you can build in `python` the following neural network and prepare the training data: ```python import torch n_tr = 7 n_val = 300 x_val = torch.linspace(-5, 5, n_val).view(-1, 1) y_val = torch.sin(x_val) x_train = torch.linspace(-3.14, 3.14, n_tr).view(-1, 1) y_train = torch.sin(x_train) + torch.randn_like(x_train) * 0.1 class Data(torch.nn.Module): def __init__(self): super(Data, self).__init__() self.register_buffer('x_val', x_val) self.register_buffer('y_val', y_val) self.register_buffer('x_train', x_train) self.register_buffer('y_train', y_train) class Net(torch.nn.Module): def __init__(self): super(Net, self).__init__() self.l1 = torch.nn.Linear(1, 10, bias = True) self.l2 = torch.nn.Linear(10, 10, bias = True) self.l3 = torch.nn.Linear(10, 1, bias = True) def forward(self, x): x = self.l1(x) x = torch.relu(x) x = self.l2(x) x = torch.relu(x) x = self.l3(x) return x class Loss(torch.nn.Module): def __init__(self, target): super(Loss, self).__init__() self.register_buffer('target', target) self.loss = torch.nn.MSELoss() def forward(self, x): return self.loss(x, self.target) # Generated TorchScript modules and serialise torch.jit.script(Data()).save('data.pt') torch.jit.script(Net()).save('net.pt') torch.jit.script(Loss(y_train)).save('loss.pt') ``` You can then load the modules into `kotlin` and train them: ```kotlin NoaFloat { // Load the serialised JIT modules // The training data val dataModule = loadJitModule("data.pt") // The DL model val netModule = loadJitModule("net.pt") // The loss function val lossModule = loadJitModule("loss.pt") // Get the tensors from the module val xTrain = dataModule.getBuffer("x_train") val yTrain = dataModule.getBuffer("y_train") val xVal = dataModule.getBuffer("x_val") val yVal = dataModule.getBuffer("y_val") // Set the model in training mode netModule.train(true) // Loss function set for training regime lossModule.setBuffer("target", yTrain) // Compute the predictions val yPred = netModule.forward(xTrain) // Compute the training error val loss = lossModule.forward(yPred) println(loss) // Set-up the Adam optimiser with learning rate 0.005 val optimiser = netModule.adamOptimiser(0.005) // Train for 250 epochs repeat(250){ // Clean gradients optimiser.zeroGrad() // Use forwardAssign to for better memory management netModule.forwardAssign(xTrain, yPred) lossModule.forwardAssign(yPred, loss) // Backward pass loss.backward() // Update model parameters optimiser.step() if(it % 50 == 0) println("Training loss: $loss") } // Finally validate the model // Compute the predictions for the validation features netModule.forwardAssign(xVal, yPred) // Set the loss for the true values lossModule.setBuffer("target", yVal) // Compute the loss on validation dataset lossModule.forwardAssign(yPred, loss) println("Validation loss: $loss") } ``` ### Custom memory management Native memory management relies on scoping with [NoaScope](src/main/kotlin/space/kscience/kmath/noa/memory/NoaScope.kt) which is readily available within an algebra context. Manual management is also possible: ```kotlin // Create a scope val scope = NoaScope() val tensor = NoaFloat(scope){ full(5f, intArrayOf(1)) }!! // the result might be null // If the computation fails resources will be freed automatically // Otherwise it's your responsibility: scope.disposeAll() // Attempts to use tensor here is undefined behaviour ``` Contributed by [Roland Grinis](https://github.com/grinisrit)