forked from kscience/kmath
testing autograd
This commit is contained in:
parent
28fac22f12
commit
0bc2c12a05
@ -43,3 +43,68 @@ 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
|
||||
|
||||
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))
|
||||
}
|
||||
```
|
||||
|
||||
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)
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
Native memory management relies on scoping
|
||||
with [NoaScope](src/main/kotlin/space/kscience/kmath/noa/memory/NoaScope.kt)
|
||||
which is readily 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
|
||||
```
|
@ -160,8 +160,8 @@ tasks["compileJava"].dependsOn(buildCpp)
|
||||
|
||||
tasks {
|
||||
withType<Test>{
|
||||
systemProperty("java.library.path", "$home/devspace/noa/cmake-build-release/kmath")
|
||||
//"$cppBuildDir/kmath")
|
||||
systemProperty("java.library.path", //"$home/devspace/noa/cmake-build-release/kmath")
|
||||
"$cppBuildDir/kmath")
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -20,3 +20,14 @@ public fun setNumThreads(numThreads: Int): Unit {
|
||||
public fun setSeed(seed: Int): Unit {
|
||||
JNoa.setSeed(seed)
|
||||
}
|
||||
|
||||
public inline fun <T, ArrayT,
|
||||
GradTensorT : NoaTensorOverField<T>,
|
||||
GradAlgebraT : NoaPartialDivisionAlgebra<T, ArrayT, GradTensorT>>
|
||||
GradAlgebraT.withGradAt(
|
||||
tensor: GradTensorT,
|
||||
block: GradAlgebraT.(GradTensorT) -> GradTensorT
|
||||
): GradTensorT {
|
||||
tensor.requiresGrad = true
|
||||
return this.block(tensor)
|
||||
}
|
||||
|
@ -81,5 +81,4 @@ class TestAlgebra {
|
||||
testingBatchedSymEig(device)
|
||||
}
|
||||
}!!
|
||||
|
||||
}
|
||||
}
|
||||
|
@ -0,0 +1,69 @@
|
||||
/*
|
||||
* Copyright 2018-2021 KMath contributors.
|
||||
* Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
|
||||
*/
|
||||
|
||||
package space.kscience.kmath.noa
|
||||
|
||||
import kotlin.test.Test
|
||||
import kotlin.test.assertTrue
|
||||
|
||||
internal fun NoaFloat.testingAutoGrad(device: Device = Device.CPU): Unit {
|
||||
setSeed(SEED)
|
||||
val dim = 3
|
||||
val tensorX = randNormal(shape = intArrayOf(dim), device = device)
|
||||
val randFeatures = randNormal(shape = intArrayOf(dim, dim), device = device)
|
||||
val tensorSigma = randFeatures + randFeatures.transpose(0, 1)
|
||||
val tensorMu = randNormal(shape = intArrayOf(dim), device = device)
|
||||
|
||||
val expressionAtX = withGradAt(tensorX) { x ->
|
||||
0.5f * (x dot (tensorSigma dot x)) + (tensorMu dot x) + 25.9f
|
||||
}
|
||||
|
||||
val gradientAtX = expressionAtX.autoGradient(tensorX, retainGraph = true)
|
||||
val hessianAtX = expressionAtX.autoHessian(tensorX)
|
||||
val expectedGradientAtX = (tensorSigma dot tensorX) + tensorMu
|
||||
|
||||
val error = (gradientAtX - expectedGradientAtX).abs().sum() +
|
||||
(hessianAtX - tensorSigma).abs().sum()
|
||||
assertTrue(error < TOLERANCE)
|
||||
}
|
||||
|
||||
internal fun NoaFloat.testingBatchedAutoGrad(device: Device = Device.CPU): Unit {
|
||||
setSeed(SEED)
|
||||
val batch = intArrayOf(2)
|
||||
val dim = 2
|
||||
val tensorX = randNormal(shape = batch + intArrayOf(1, dim), device = device)
|
||||
val randFeatures = randNormal(shape = batch + intArrayOf(dim, dim), device = device)
|
||||
val tensorSigma = randFeatures + randFeatures.transpose(-2, -1)
|
||||
val tensorMu = randNormal(shape = batch + intArrayOf(1, dim), device = device)
|
||||
|
||||
val expressionAtX = withGradAt(tensorX) { x ->
|
||||
val xt = x.transpose(-1, -2)
|
||||
(0.5f * (x dot (tensorSigma dot xt)) + (tensorMu dot xt) + 58.2f).sumAll()
|
||||
}
|
||||
|
||||
|
||||
val gradientAtX = expressionAtX.autoGradient(tensorX)
|
||||
val expectedGradientAtX = (tensorX dot tensorSigma) + tensorMu
|
||||
|
||||
val error = (gradientAtX - expectedGradientAtX).abs().sum()
|
||||
assertTrue(error < TOLERANCE)
|
||||
}
|
||||
|
||||
class TestAutoGrad {
|
||||
|
||||
@Test
|
||||
fun testAutoGrad() = NoaFloat {
|
||||
withCuda { device ->
|
||||
testingAutoGrad(device)
|
||||
}
|
||||
}!!
|
||||
|
||||
@Test
|
||||
fun testBatchedAutoGrad() = NoaFloat {
|
||||
withCuda { device ->
|
||||
testingBatchedAutoGrad(device)
|
||||
}
|
||||
}!!
|
||||
}
|
@ -5,6 +5,7 @@
|
||||
|
||||
package space.kscience.kmath.noa
|
||||
|
||||
import space.kscience.kmath.noa.memory.NoaScope
|
||||
import kotlin.test.Test
|
||||
import kotlin.test.assertEquals
|
||||
|
||||
@ -59,4 +60,16 @@ class TestUtils {
|
||||
}
|
||||
}!!
|
||||
|
||||
@Test
|
||||
fun testScoping(): Unit {
|
||||
val scope = NoaScope()
|
||||
val tensor = NoaFloat(scope){
|
||||
full(5f, intArrayOf(1))
|
||||
}!!
|
||||
assertEquals(tensor.numElements, 1)
|
||||
assertEquals(scope.disposables.size, 1)
|
||||
scope.disposeAll()
|
||||
assertEquals(scope.disposables.size, 0)
|
||||
}
|
||||
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user