From 6c4741ede6558f53e5b90e35f934ba0a83ceb13a Mon Sep 17 00:00:00 2001 From: Alexander Nozik Date: Wed, 20 Oct 2021 16:06:45 +0300 Subject: [PATCH] [WIP] TensorFlow --- kmath-tensorflow/build.gradle.kts | 14 ++ .../kmath/tensorflow/TensorFlowAlgebra.kt | 220 ++++++++++++++++++ settings.gradle.kts | 1 + 3 files changed, 235 insertions(+) create mode 100644 kmath-tensorflow/build.gradle.kts create mode 100644 kmath-tensorflow/src/main/kotlin/space/kscience/kmath/tensorflow/TensorFlowAlgebra.kt diff --git a/kmath-tensorflow/build.gradle.kts b/kmath-tensorflow/build.gradle.kts new file mode 100644 index 000000000..fa77a272a --- /dev/null +++ b/kmath-tensorflow/build.gradle.kts @@ -0,0 +1,14 @@ +plugins { + id("ru.mipt.npm.gradle.jvm") +} + +description = "Google tensorflow connector" + +dependencies { + api(project(":kmath-tensors")) + api("org.tensorflow:tensorflow-core-api:0.3.3") +} + +readme { + maturity = ru.mipt.npm.gradle.Maturity.PROTOTYPE +} \ No newline at end of file diff --git a/kmath-tensorflow/src/main/kotlin/space/kscience/kmath/tensorflow/TensorFlowAlgebra.kt b/kmath-tensorflow/src/main/kotlin/space/kscience/kmath/tensorflow/TensorFlowAlgebra.kt new file mode 100644 index 000000000..849a2e909 --- /dev/null +++ b/kmath-tensorflow/src/main/kotlin/space/kscience/kmath/tensorflow/TensorFlowAlgebra.kt @@ -0,0 +1,220 @@ +package space.kscience.kmath.tensorflow + + +import org.tensorflow.Graph +import org.tensorflow.Operand +import org.tensorflow.Output +import org.tensorflow.Session +import org.tensorflow.ndarray.NdArray +import org.tensorflow.op.Ops +import org.tensorflow.op.core.Constant +import org.tensorflow.types.family.TType +import space.kscience.kmath.misc.PerformancePitfall +import space.kscience.kmath.nd.Shape +import space.kscience.kmath.tensors.api.Tensor +import space.kscience.kmath.tensors.api.TensorAlgebra + +private fun IntArray.toLongArray() = LongArray(size) { get(it).toLong() } +private fun LongArray.toIntArray() = IntArray(size) { get(it).toInt() } + +private val NdArray.scalar: T + get() = getObject() + + +public sealed interface TensorFlowTensor : Tensor + +@JvmInline +public value class TensorFlowArray(public val tensor: NdArray) : Tensor { + override val shape: Shape get() = tensor.shape().asArray().toIntArray() + + override fun get(index: IntArray): T = tensor.getObject(*index.toLongArray()) + + @PerformancePitfall + override fun elements(): Sequence> = sequence { + tensor.scalars().forEachIndexed { index: LongArray, ndArray: NdArray -> + //yield(index.toIntArray() to ndArray.scalar) + TODO() + } + } + + override fun set(index: IntArray, value: T) { + tensor.setObject(value, *index.toLongArray()) + } +} + +public abstract class TensorFlowOutput( + private val graph: Graph, + output: Output +) : TensorFlowTensor { + + public var output: Output = output + internal set + + override val shape: Shape get() = output.shape().asArray().toIntArray() + + protected abstract fun org.tensorflow.Tensor.actualizeTensor(): NdArray + + private val actualTensor by lazy { + val session = Session(graph) + TensorFlowArray(session.runner().fetch(output).run().first().actualizeTensor()) + } + + override fun get(index: IntArray): T = actualTensor[index] + + @PerformancePitfall + override fun elements(): Sequence> = actualTensor.elements() + + override fun set(index: IntArray, value: T) { + actualTensor[index] = value + } + +} + + +public abstract class TensorFlowAlgebra internal constructor( + private val graph: Graph +) : TensorAlgebra { + + private val ops by lazy { Ops.create(graph) } + + protected fun Tensor.asTensorFlow(): TensorFlowOutput = if (this is TensorFlowOutput) this else { + TODO() + } + + protected abstract fun Output.wrap(): TensorFlowOutput + + protected abstract fun const(value: T): Constant + + override fun Tensor.valueOrNull(): T? = if (shape contentEquals intArrayOf(1)) + get(Shape(0)) else null + + private inline fun Tensor.biOp( + other: Tensor, + operation: (left: Operand, right: Operand) -> Operand + ): TensorFlowOutput { + val left = asTensorFlow().output + val right = other.asTensorFlow().output + return operation(left, right).asOutput().wrap() + } + + private inline fun T.biOp( + other: Tensor, + operation: (left: Operand, right: Operand) -> Operand + ): TensorFlowOutput { + val left = const(this) + val right = other.asTensorFlow().output + return operation(left, right).asOutput().wrap() + } + + private inline fun Tensor.biOp( + value: T, + operation: (left: Operand, right: Operand) -> Operand + ): TensorFlowOutput { + val left = asTensorFlow().output + val right = const(value) + return operation(left, right).asOutput().wrap() + } + + private inline fun Tensor.inPlaceOp( + other: Tensor, + operation: (left: Operand, right: Operand) -> Operand + ): Unit { + val origin = asTensorFlow() + val left = origin.output + val right = other.asTensorFlow().output + origin.output = operation(left, right).asOutput() + } + + private inline fun Tensor.inPlaceOp( + value: T, + operation: (left: Operand, right: Operand) -> Operand + ): Unit { + val origin = asTensorFlow() + val left = origin.output + val right = const(value) + origin.output = operation(left, right).asOutput() + } + + private inline fun unOp(value: Tensor, operation: (Operand) -> Operand): TensorFlowOutput = + operation(value.asTensorFlow().output).asOutput().wrap() + + override fun T.plus(other: Tensor) = biOp(other, ops.math::add) + + override fun Tensor.plus(value: T) = biOp(value, ops.math::add) + + override fun Tensor.plus(other: Tensor) = biOp(other, ops.math::add) + + override fun Tensor.plusAssign(value: T): Unit = inPlaceOp(value, ops.math::add) + + override fun Tensor.plusAssign(other: Tensor): Unit = inPlaceOp(other, ops.math::add) + + override fun Tensor.minus(value: T) = biOp(value, ops.math::sub) + + override fun Tensor.minus(other: Tensor) = biOp(other, ops.math::sub) + + override fun Tensor.minusAssign(value: T): Unit = inPlaceOp(value, ops.math::sub) + + override fun Tensor.minusAssign(other: Tensor): Unit = inPlaceOp(other, ops.math::sub) + + override fun T.times(other: Tensor) = biOp(other, ops.math::mul) + + override fun Tensor.times(value: T) = biOp(value, ops.math::mul) + + override fun Tensor.times(other: Tensor): TensorFlowOutput = biOp(other, ops.math::mul) + + override fun Tensor.timesAssign(value: T): Unit = inPlaceOp(value, ops.math::mul) + + override fun Tensor.timesAssign(other: Tensor): Unit = inPlaceOp(other, ops.math::mul) + + override fun Tensor.unaryMinus() = unOp(this, ops.math::neg) + + override fun Tensor.get(i: Int): Tensor{ + ops. + } + + override fun Tensor.transpose(i: Int, j: Int): Tensor { + TODO("Not yet implemented") + } + + override fun Tensor.view(shape: IntArray): Tensor { + TODO("Not yet implemented") + } + + override fun Tensor.viewAs(other: Tensor): Tensor { + TODO("Not yet implemented") + } + + override fun Tensor.dot(other: Tensor) = biOp(other, ops.math.) + + override fun diagonalEmbedding(diagonalEntries: Tensor, offset: Int, dim1: Int, dim2: Int): Tensor = ops.run { + TODO("Not yet implemented") + } + + override fun Tensor.sum(): T { + TODO("Not yet implemented") + } + + override fun Tensor.sum(dim: Int, keepDim: Boolean): Tensor { + TODO("Not yet implemented") + } + + override fun Tensor.min(): T { + TODO("Not yet implemented") + } + + override fun Tensor.min(dim: Int, keepDim: Boolean): Tensor { + TODO("Not yet implemented") + } + + override fun Tensor.max(): T { + TODO("Not yet implemented") + } + + override fun Tensor.max(dim: Int, keepDim: Boolean): Tensor { + TODO("Not yet implemented") + } + + override fun Tensor.argMax(dim: Int, keepDim: Boolean): Tensor { + TODO("Not yet implemented") + } +} \ No newline at end of file diff --git a/settings.gradle.kts b/settings.gradle.kts index e73381bf2..a29f33824 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -33,6 +33,7 @@ include( ":kmath-commons", ":kmath-viktor", ":kmath-multik", + ":kmath-tensorflow", ":kmath-optimization", ":kmath-stat", ":kmath-nd4j",