@hoff97/tensor-js
Version:
PyTorch like deep learning inferrence library
476 lines • 21.5 kB
JavaScript
import Tensor, { tensorValuesConstructor, } from '../../types';
import { compareShapes, getSize } from '../../util/shape';
import { MatMulOperation } from '../../ops/gpu/matMul/matmul';
import { defaultAllocator, gl } from './gl';
import { ExpOperation } from '../../ops/gpu/unary/exp';
import { ConvBiasOperation, ConvOperation } from '../../ops/gpu/conv/conv';
import { AbsOperation } from '../../ops/gpu/unary/abs';
import { AddOperation } from '../../ops/gpu/binary/add';
import { MultiplyOperation } from '../../ops/gpu/binary/multiply';
import { SubtractOperation } from '../../ops/gpu/binary/subtract';
import { DivideOperation } from '../../ops/gpu/binary/divide';
import { AveragePoolOperation } from '../../ops/gpu/conv/averagePool';
import { ReduceMeanOperation } from '../../ops/gpu/pool/reduceMean';
import { ReduceMeanSquareOperation } from '../../ops/gpu/pool/reduceMeanSquare';
import { SumSquareOperation } from '../../ops/gpu/pool/sumSquare';
import { SumOperation } from '../../ops/gpu/pool/sum';
import { ProductOperation } from '../../ops/gpu/pool/product';
import { MaxOperation } from '../../ops/gpu/pool/max';
import { MinOperation } from '../../ops/gpu/pool/min';
import { CeilOperation } from '../../ops/gpu/unary/ceil';
import { ClipOperation } from '../../ops/gpu/unary/clip';
import { FloorOperation } from '../../ops/gpu/unary/floor';
import { ConcatOperation } from '../../ops/gpu/util/concat';
import { CopyOperation } from '../../ops/gpu/util/copy';
import { ExpandOperation } from '../../ops/gpu/util/expand';
import { GatherOperation } from '../../ops/gpu/util/gather';
import { GemmCOperation, GemmOperation } from '../../ops/gpu/matMul/gemm';
import { PowerOperation } from '../../ops/gpu/binary/power';
import { SqrtOperation } from '../../ops/gpu/unary/sqrt';
import { LogOperation } from '../../ops/gpu/unary/log';
import { TransposeOperation } from '../../ops/gpu/util/transpose';
import { RepeatOperation } from '../../ops/gpu/util/repeat';
import { PadOperation } from '../../ops/gpu/conv/pad';
import { SliceOperation } from '../../ops/gpu/util/slice';
import { UpsampleOperation } from '../../ops/gpu/conv/upsample';
import { NormalizeOperation } from '../../ops/gpu/conv/normalize';
import { Dispatcher } from '../../ops/gpu/dispatcher';
import { SignOperation } from '../../ops/gpu/unary/sign';
import { NegateOperation } from '../../ops/gpu/unary/negate';
import { ClipBackwardOperation } from '../../ops/gpu/util/clipBackward';
import { ConvTransposeOperation } from '../../ops/gpu/conv/convTranspose';
import { SigmoidOperation } from '../../ops/gpu/unary/sigmoid';
import { AddMultiplyScalarOperation } from '../../ops/gpu/unary/addMultiplyScalar';
import { SetValuesOperation } from '../../ops/gpu/util/setValues';
import { ASinOperation, SinHOperation, SinOperation, } from '../../ops/gpu/unary/sin';
import { ACosOperation, CosHOperation, CosOperation, } from '../../ops/gpu/unary/cos';
import { ATanOperation, TanHOperation, TanOperation, } from '../../ops/gpu/unary/tan';
import { ReduceLogSumOperation } from '../../ops/gpu/pool/reduceLogSum';
import { ReduceLogSumExpOperation } from '../../ops/gpu/pool/reduceLogSumExp';
import { HardSigmoidOperation } from '../../ops/gpu/unary/hardSigmoid';
import { PowerScalarOperation } from '../../ops/gpu/unary/powerScalar';
import { RoundOperation } from '../../ops/gpu/unary/round';
export class GPUTensor extends Tensor {
constructor(values, shape, dtype) {
super(dtype || 'float32');
this.shape = shape;
this.deleted = false;
this.size = getSize(shape);
if (values instanceof Array) {
this.memory = defaultAllocator.allocateTexture(values, this.dtype);
}
else {
this.memory = values;
}
}
static range(start, limit, delta, dtype) {
const size = Math.max(Math.ceil((limit - start) / delta), 0);
const values = new Array(size);
for (let i = 0; i < size; i++) {
values[i] = start + i * delta;
}
return new GPUTensor(values, [size], dtype);
}
static fromData(data) {
const texture = gl.texture({
data: data,
format: 'rgba',
type: defaultAllocator.getColorType('float32'),
});
const memory = defaultAllocator.allocateFramebuffer(texture, 'float32');
const width = texture.width;
const height = texture.height;
return new GPUTensor(memory, [height, width, 4], 'float32');
}
cast(dtype) {
if (dtype === 'float64') {
throw new Error('The WebGL backend does not support float64 tensors');
}
return defaultCopyD.calc({ input: this }, dtype);
}
getValues() {
if (this.dtype !== 'float16') {
return new Promise(resolve => {
gl({ framebuffer: this.memory.frameBuffer })(() => {
let result = new Float32Array(this.memory.size);
result = gl.read(result);
if (this.dtype === 'float32') {
resolve(result.subarray(0, this.size));
}
else {
const arr = new tensorValuesConstructor[this.dtype](this.size);
for (let i = 0; i < this.size; i++) {
arr[i] = result[i];
}
resolve(arr);
}
});
});
}
throw new Error('Reading values not supported for data type float16');
}
getShape() {
return this.shape;
}
constantLike(value) {
return new GPUTensor(new Array(this.size).fill(value), this.shape, this.dtype);
}
singleConstant(value) {
return new GPUTensor([value], [1], this.dtype);
}
delete() {
if (!this.deleted) {
this.deleted = true;
defaultAllocator.deallocate(this.memory);
//@ts-ignore
this.memory = undefined;
}
}
copy() {
return defaultCopyD.calc({ input: this }, this.dtype);
}
exp() {
return defaultExpD.calc({ input: this }, this.dtype);
}
log() {
return defaultLogD.calc({ input: this }, this.dtype);
}
sqrt() {
return defaultSqrtD.calc({ input: this }, this.dtype);
}
abs() {
return defaultAbsD.calc({ input: this }, this.dtype);
}
sin() {
return defaultSinD.calc({ input: this }, this.dtype);
}
cos() {
return defaultCosD.calc({ input: this }, this.dtype);
}
tan() {
return defaultTanD.calc({ input: this }, this.dtype);
}
asin() {
return defaultASinD.calc({ input: this }, this.dtype);
}
acos() {
return defaultACosD.calc({ input: this }, this.dtype);
}
atan() {
return defaultATanD.calc({ input: this }, this.dtype);
}
sinh() {
return defaultSinHD.calc({ input: this }, this.dtype);
}
cosh() {
return defaultCosHD.calc({ input: this }, this.dtype);
}
tanh() {
return defaultTanHD.calc({ input: this }, this.dtype);
}
asinh() {
throw new Error('Method not implemented');
}
acosh() {
throw new Error('Method not implemented');
}
atanh() {
throw new Error('Method not implemented');
}
sigmoid() {
return defaultSigmoidD.calc({ input: this }, this.dtype);
}
hardSigmoid(alpha, beta) {
return defaultHardSigmoidD.calc({ input: this, alpha, beta }, this.dtype);
}
floor() {
return defaultFloorD.calc({ input: this }, this.dtype);
}
ceil() {
return defaultCeilD.calc({ input: this }, this.dtype);
}
round() {
return defaultRoundD.calc({ input: this }, this.dtype);
}
negate() {
return defaultNegateD.calc({ input: this }, this.dtype);
}
addMultiplyScalar(factor, add) {
return defaultAddMultiplyScalarD.calc({ input: this, factor, add }, this.dtype);
}
powerScalar(power, factor) {
return defaultPowerScalarD.calc({ input: this, factor, power }, this.dtype);
}
sign() {
return defaultSignD.calc({ input: this }, this.dtype);
}
setValues(values, starts) {
if (!(values instanceof GPUTensor)) {
throw new Error('Can only set GPU values to GPU tensor');
}
return defaultSetValuesD.calc({ A: this, Values: values, starts }, this.dtype);
}
add_impl(th, tensor, resultShape, alpha, beta) {
if (!(tensor instanceof GPUTensor) || !(th instanceof GPUTensor)) {
throw new Error('Can only add GPU tensor to GPU tensor');
}
return defaultAddD.calc({ A: th, B: tensor, outputShape: resultShape, alpha, beta }, this.dtype);
}
subtract_impl(th, tensor, resultShape, alpha, beta) {
if (!(tensor instanceof GPUTensor) || !(th instanceof GPUTensor)) {
throw new Error('Can only subtract GPU tensor from GPU tensor');
}
return defaultSubtractD.calc({ A: th, B: tensor, outputShape: resultShape, alpha, beta }, this.dtype);
}
multiply_impl(th, tensor, resultShape, alpha) {
if (!(tensor instanceof GPUTensor) || !(th instanceof GPUTensor)) {
throw new Error('Can only multiply GPU tensor with GPU tensor');
}
return defaultMultiplyD.calc({ A: th, B: tensor, outputShape: resultShape, alpha }, this.dtype);
}
divide_impl(th, tensor, resultShape, alpha) {
if (!(tensor instanceof GPUTensor) || !(th instanceof GPUTensor)) {
throw new Error('Can only divide GPU tensor by GPU tensor');
}
return defaultDivideD.calc({ A: th, B: tensor, outputShape: resultShape, alpha }, this.dtype);
}
power_impl(th, tensor, resultShape) {
if (!(tensor instanceof GPUTensor) || !(th instanceof GPUTensor)) {
throw new Error('Can only take GPU tensor to power of GPU tensor');
}
return defaultPowerD.calc({ A: th, B: tensor, outputShape: resultShape }, this.dtype);
}
matMul(tensor) {
if (!(tensor instanceof GPUTensor)) {
throw new Error('Can only matrix multiply GPU tensor to GPU tensor');
}
return defaultMatMulD.calc({ A: this, B: tensor }, this.dtype);
}
gemm_impl(b, aTranspose, bTranspose, alpha, beta, c) {
if (!(b instanceof GPUTensor && (c === undefined || c instanceof GPUTensor))) {
throw new Error('Can only do gemm with CPU tensors');
}
if (c === undefined) {
return defaultGemmD.calc({ a: this, b, aTranspose, bTranspose, alpha, beta }, this.dtype);
}
else {
return defaultGemmCD.calc({
a: this,
b,
c: c,
aTranspose,
bTranspose,
alpha,
beta,
}, this.dtype);
}
}
sum_impl(axes, keepDims) {
return defaultSumD.calc({ X: this, axes, keepDims }, this.dtype);
}
sumSquare_impl(axes, keepDims) {
return defaultSumSquareD.calc({ X: this, axes, keepDims }, this.dtype);
}
reduceMean_impl(axes, keepDims) {
return defaultMeanD.calc({ X: this, axes, keepDims }, this.dtype);
}
reduceMeanSquare_impl(axes, keepDims) {
return defaultMeanSquareD.calc({ X: this, axes, keepDims }, this.dtype);
}
reduceLogSum_impl(axes, keepDims) {
return defaultLogSumD.calc({ X: this, axes, keepDims }, this.dtype);
}
reduceLogSumExp_impl(axes, keepDims) {
return defaultLogSumExpD.calc({ X: this, axes, keepDims }, this.dtype);
}
product_impl(axes, keepDims) {
return defaultProductD.calc({ X: this, axes, keepDims }, this.dtype);
}
max_impl(axes, keepDims) {
return defaultMaxD.calc({ X: this, axes, keepDims }, this.dtype);
}
min_impl(axes, keepDims) {
return defaultMinD.calc({ X: this, axes, keepDims }, this.dtype);
}
conv_impl(kernel, dilations, group, pads, strides, activation, bias) {
if (!(kernel instanceof GPUTensor) ||
(bias !== undefined && !(bias instanceof GPUTensor))) {
throw new Error('Can only do convolution of GPU tensor with GPU tensor');
}
if (bias === undefined) {
return defaultConvD.calc({
X: this,
W: kernel,
pads,
dilations,
strides,
activation,
}, this.dtype);
}
else {
return defaultConvBiasD.calc({
X: this,
W: kernel,
B: bias,
pads,
dilations,
strides,
activation,
}, this.dtype);
}
}
convTranspose_impl(kernel, dilations, group, pads, strides) {
if (!(kernel instanceof GPUTensor)) {
throw new Error('Can only do transpose convolution of GPU tensor with GPU tensor');
}
return defaultConvTransposeD.calc({
X: this,
W: kernel,
pads,
dilations,
strides,
}, this.dtype);
}
averagePool_impl(kernelShape, pads, strides, includePad) {
return defaultAveragePoolD.calc({
X: this,
includePad,
kernelShape,
pads,
strides,
}, this.dtype);
}
reshape_impl(shape, _copy) {
if (_copy) {
return defaultCopyD.calc({ input: this, outputShape: shape }, this.dtype);
}
else {
return new GPUTensor(this.memory, shape, this.dtype);
}
}
concat(tensor, axis) {
if (!(tensor instanceof GPUTensor)) {
throw new Error('Can only concat GPU tensor to GPU tensor');
}
if (axis < 0) {
axis += this.shape.length;
}
return defaultConcatD.calc({ A: this, B: tensor, axis }, this.dtype);
}
transpose_impl(permutation) {
return defaultTransposeD.calc({ A: this, permutation }, this.dtype);
}
clip(min, max) {
return defaultClipD.calc({ input: this, minVal: min, maxVal: max }, this.dtype);
}
clipBackward(grad, min, max) {
return defaultClipBackwardD.calc({ input: this, minVal: min, maxVal: max, grad }, this.dtype);
}
repeat(repeats) {
return defaultRepeatD.calc({ A: this, repeats }, this.dtype);
}
expand(shape) {
// eslint-disable-next-line @typescript-eslint/no-unused-vars
const [_shape, _o, resultShape] = this.alignShapes(this.shape, shape);
if (compareShapes(this.shape, resultShape)) {
return this.copy();
}
return defaultExpandD.calc({
input: this.reshape(_shape, false),
outputShape: resultShape,
}, this.dtype);
}
pad_impl(pads, mode, value) {
return defaultPadD.calc({ input: this, pads, mode, value }, this.dtype);
}
gather(axis, indices) {
return defaultGatherD.calc({ X: this, axis, indices }, this.dtype);
}
slice_impl(starts, ends, axes, steps) {
return defaultSliceD.calc({ X: this, starts, ends, axes, steps }, this.dtype);
}
upsample(scales) {
return defaultUpsampleD.calc({ X: this, scales }, this.dtype);
}
normalize(mean, variance, epsilon, scale, bias) {
if (!(mean instanceof GPUTensor) ||
!(variance instanceof GPUTensor) ||
!(scale instanceof GPUTensor) ||
!(bias instanceof GPUTensor)) {
throw new Error('Can only normalize with CPU tensors');
}
return defaultNormalizeD.calc({
X: this,
Mean: mean,
Variance: variance,
Scale: scale,
Bias: bias,
epsilon,
}, this.dtype);
}
}
export function gpuConstructor(a, b, dtype) {
return new GPUTensor(a, b, dtype);
}
const defaultMatMulD = new Dispatcher((dtype) => new MatMulOperation(gpuConstructor, dtype));
const defaultGemmD = new Dispatcher((dtype) => new GemmOperation(gpuConstructor, dtype));
const defaultGemmCD = new Dispatcher((dtype) => new GemmCOperation(gpuConstructor, dtype));
//Unary operations
const defaultExpD = new Dispatcher((dtype) => new ExpOperation(gpuConstructor, dtype));
const defaultAbsD = new Dispatcher((dtype) => new AbsOperation(gpuConstructor, dtype));
const defaultSinD = new Dispatcher((dtype) => new SinOperation(gpuConstructor, dtype));
const defaultCosD = new Dispatcher((dtype) => new CosOperation(gpuConstructor, dtype));
const defaultTanD = new Dispatcher((dtype) => new TanOperation(gpuConstructor, dtype));
const defaultASinD = new Dispatcher((dtype) => new ASinOperation(gpuConstructor, dtype));
const defaultACosD = new Dispatcher((dtype) => new ACosOperation(gpuConstructor, dtype));
const defaultATanD = new Dispatcher((dtype) => new ATanOperation(gpuConstructor, dtype));
const defaultSinHD = new Dispatcher((dtype) => new SinHOperation(gpuConstructor, dtype));
const defaultCosHD = new Dispatcher((dtype) => new CosHOperation(gpuConstructor, dtype));
const defaultTanHD = new Dispatcher((dtype) => new TanHOperation(gpuConstructor, dtype));
const defaultSigmoidD = new Dispatcher((dtype) => new SigmoidOperation(gpuConstructor, dtype));
const defaultHardSigmoidD = new Dispatcher((dtype) => new HardSigmoidOperation(gpuConstructor, dtype));
const defaultCeilD = new Dispatcher((dtype) => new CeilOperation(gpuConstructor, dtype));
const defaultFloorD = new Dispatcher((dtype) => new FloorOperation(gpuConstructor, dtype));
const defaultRoundD = new Dispatcher((dtype) => new RoundOperation(gpuConstructor, dtype));
const defaultClipD = new Dispatcher((dtype) => new ClipOperation(gpuConstructor, dtype));
const defaultClipBackwardD = new Dispatcher((dtype) => new ClipBackwardOperation(gpuConstructor, dtype));
const defaultSqrtD = new Dispatcher((dtype) => new SqrtOperation(gpuConstructor, dtype));
const defaultLogD = new Dispatcher((dtype) => new LogOperation(gpuConstructor, dtype));
const defaultNegateD = new Dispatcher((dtype) => new NegateOperation(gpuConstructor, dtype));
const defaultAddMultiplyScalarD = new Dispatcher((dtype) => new AddMultiplyScalarOperation(gpuConstructor, dtype));
const defaultPowerScalarD = new Dispatcher((dtype) => new PowerScalarOperation(gpuConstructor, dtype));
const defaultSignD = new Dispatcher((dtype) => new SignOperation(gpuConstructor, dtype));
//Convolutions
const defaultConvD = new Dispatcher((dtype) => new ConvOperation(gpuConstructor, dtype));
const defaultAveragePoolD = new Dispatcher((dtype) => new AveragePoolOperation(gpuConstructor, dtype));
const defaultConvBiasD = new Dispatcher((dtype) => new ConvBiasOperation(gpuConstructor, dtype));
const defaultConvTransposeD = new Dispatcher((dtype) => new ConvTransposeOperation(gpuConstructor, dtype));
const defaultPadD = new Dispatcher((dtype) => new PadOperation(gpuConstructor, dtype));
const defaultUpsampleD = new Dispatcher((dtype) => new UpsampleOperation(gpuConstructor, dtype));
//Binary operations
const defaultAddD = new Dispatcher((dtype) => new AddOperation(gpuConstructor, dtype));
const defaultSubtractD = new Dispatcher((dtype) => new SubtractOperation(gpuConstructor, dtype));
const defaultMultiplyD = new Dispatcher((dtype) => new MultiplyOperation(gpuConstructor, dtype));
const defaultDivideD = new Dispatcher((dtype) => new DivideOperation(gpuConstructor, dtype));
const defaultPowerD = new Dispatcher((dtype) => new PowerOperation(gpuConstructor, dtype));
//Reductions
const defaultMeanD = new Dispatcher((dtype) => new ReduceMeanOperation(gpuConstructor, dtype));
const defaultMeanSquareD = new Dispatcher((dtype) => new ReduceMeanSquareOperation(gpuConstructor, dtype));
const defaultSumSquareD = new Dispatcher((dtype) => new SumSquareOperation(gpuConstructor, dtype));
const defaultSumD = new Dispatcher((dtype) => new SumOperation(gpuConstructor, dtype));
const defaultProductD = new Dispatcher((dtype) => new ProductOperation(gpuConstructor, dtype));
const defaultMaxD = new Dispatcher((dtype) => new MaxOperation(gpuConstructor, dtype));
const defaultMinD = new Dispatcher((dtype) => new MinOperation(gpuConstructor, dtype));
const defaultLogSumD = new Dispatcher((dtype) => new ReduceLogSumOperation(gpuConstructor, dtype));
const defaultLogSumExpD = new Dispatcher((dtype) => new ReduceLogSumExpOperation(gpuConstructor, dtype));
//Util
const defaultConcatD = new Dispatcher((dtype) => new ConcatOperation(gpuConstructor, dtype));
const defaultSetValuesD = new Dispatcher((dtype) => new SetValuesOperation(gpuConstructor, dtype));
const defaultCopyD = new Dispatcher((dtype) => new CopyOperation(gpuConstructor, dtype));
const defaultExpandD = new Dispatcher((dtype) => new ExpandOperation(gpuConstructor, dtype));
const defaultGatherD = new Dispatcher((dtype) => new GatherOperation(gpuConstructor, dtype));
const defaultTransposeD = new Dispatcher((dtype) => new TransposeOperation(gpuConstructor, dtype));
const defaultRepeatD = new Dispatcher((dtype) => new RepeatOperation(gpuConstructor, dtype));
const defaultSliceD = new Dispatcher((dtype) => new SliceOperation(gpuConstructor, dtype));
const defaultNormalizeD = new Dispatcher((dtype) => new NormalizeOperation(gpuConstructor, dtype));
//# sourceMappingURL=tensor.js.map