Skip to content
This repository has been archived by the owner on Nov 16, 2023. It is now read-only.

Topk-1 CPU #300

Closed
wants to merge 2 commits into from
Closed

Topk-1 CPU #300

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8ec5828
Fixes / more tests for TopK-1
J3698 May 22, 2021
5147243
Fix TopK bug
J3698 May 25, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .gitignore
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,3 +17,4 @@
npm-debug.log
.DS_Store
yarn-error.log
*.swp
2 changes: 1 addition & 1 deletion docs/operators.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,7 @@ _This file is automatically generated from the def files via [this script](/tool
| [TfIdfVectorizer](https://github.com/onnx/onnx/blob/master/docs/Operators.md#TfIdfVectorizer) | | | |
| [ThresholdedRelu](https://github.com/onnx/onnx/blob/master/docs/Operators.md#ThresholdedRelu) | | | |
| [Tile](https://github.com/onnx/onnx/blob/master/docs/Operators.md#Tile) | [6+](https://github.com/onnx/onnx/blob/master/docs/Changelog.md#Tile-6) | | [6+](https://github.com/onnx/onnx/blob/master/docs/Changelog.md#Tile-6) |
| [TopK](https://github.com/onnx/onnx/blob/master/docs/Operators.md#TopK) | | | |
| [TopK](https://github.com/onnx/onnx/blob/master/docs/Operators.md#TopK) | [1-9](https://github.com/onnx/onnx/blob/master/docs/Changelog.md#TopK-1) | | |
| [Transpose](https://github.com/onnx/onnx/blob/master/docs/Operators.md#Transpose) | [1+](https://github.com/onnx/onnx/blob/master/docs/Changelog.md#Transpose-1) | | [1+](https://github.com/onnx/onnx/blob/master/docs/Changelog.md#Transpose-1) |
| [Unique](https://github.com/onnx/onnx/blob/master/docs/Operators.md#Unique) | | | |
| [Unsqueeze](https://github.com/onnx/onnx/blob/master/docs/Operators.md#Unsqueeze) | [1-10](https://github.com/onnx/onnx/blob/master/docs/Changelog.md#Unsqueeze-1), [11+](https://github.com/onnx/onnx/blob/master/docs/Changelog.md#Unsqueeze-11) | | [1-10](https://github.com/onnx/onnx/blob/master/docs/Changelog.md#Unsqueeze-1), [11+](https://github.com/onnx/onnx/blob/master/docs/Changelog.md#Unsqueeze-11) |
Expand Down
2 changes: 2 additions & 0 deletions lib/backends/cpu/op-resolve-rules.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ import {CpuSoftmax} from './ops/softmax';
import {CpuSqueeze} from './ops/squeeze';
import {CpuSum} from './ops/sum';
import {CpuTile} from './ops/tile';
import {CpuTopK} from './ops/topK';
import {CpuTranspose} from './ops/transpose';
import * as unaryOps from './ops/unary-op';
import {CpuUnaryOp} from './ops/unary-op';
Expand Down Expand Up @@ -108,6 +109,7 @@ export const CPU_OP_RESOLVE_RULES: ReadonlyArray<OpSet.ResolveRule> = [
['Tan', '', '7+', () => new CpuUnaryOp(FLOAT_TYPES, unaryOps.tan)],
['Tanh', '', '6+', () => new CpuUnaryOp(FLOAT_TYPES, unaryOps.tanh)],
['Tile', '', '6+', () => new CpuTile()],
['TopK', '', '1-9', () => new CpuTopK()],
['Transpose', '', '1+', () => new CpuTranspose()],
['Unsqueeze', '', '1+', () => new CpuUnsqueeze()],
['Upsample', '', '7-8', () => new CpuUpsample(7)],
Expand Down
85 changes: 85 additions & 0 deletions lib/backends/cpu/ops/topK.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,85 @@
import {TopK} from '../../../ops/topK';
import {Tensor} from '../../../tensor';
import {assert, ShapeUtil} from '../../../util';
import {CpuInferenceHandler} from '../inference-handler';

export class CpuTopK extends TopK {
run(inferenceHandler: CpuInferenceHandler, inputs: Tensor[]): Tensor[] {
const output = tensorTopK(inputs[0], this.k, this.axis);
return output;
}
}

export function tensorTopK(x: Tensor, K: number, axis: number): Tensor[] {
assert(K === Math.round(K), () => 'K should be an integer');
K = Math.round(K);

const rank = x.dims ? x.dims.length : 1;
axis = ShapeUtil.normalizeAxis(axis, rank);
const outputDims = [...x.dims];
outputDims[axis] = K;
const dimsWithoutAxis = [...x.dims];
dimsWithoutAxis[axis] = 1;

const values = new Tensor(outputDims, x.type);
const indices = new Tensor(outputDims, 'int32');

const blockSize = ShapeUtil.sizeFromDimension(x.dims, axis + 1);
const array = Array<number>(x.dims[axis]);
const elems = ShapeUtil.size(outputDims);
const xStrides = ShapeUtil.computeStrides(x.dims);
const withoutAxisStrides = ShapeUtil.computeStrides(dimsWithoutAxis);
const outputStrides = ShapeUtil.computeStrides(outputDims);
for (let i = 0; i < elems; ++i) {
const startIdx = ShapeUtil.offsetToIndices(i, withoutAxisStrides);
const offset = ShapeUtil.indicesToOffset(startIdx, xStrides);
for (let j = 0; j < x.dims[axis]; ++j) {
array[j] = x.numberData[offset + j * blockSize];
}

const topk = listTopK(array, K);
const dstOffset = ShapeUtil.indicesToOffset(startIdx, outputStrides);
for (let j = 0; j < K; ++j) {
values.numberData[dstOffset + j * blockSize] = topk[j][0];
indices.numberData[dstOffset + j * blockSize] = topk[j][1];
}
}

return [values, indices];
}

export function listTopK(list: number[], K: number): Array<[number, number]> {
const enumList: Array<[number, number]> = list.map((el, idx) => [el, idx]);
let start = 0;
let end = list.length;

let iters = 0;
while (start !== K && iters <= list.length) {
let curr = start;
for (let i = start; i < end; ++i) {
if (enumList[i][0] > enumList[end - 1][0] ||
(enumList[i][0] === enumList[end - 1][0] && enumList[i][1] <= enumList[end - 1][1])) {
[enumList[i], enumList[curr]] = [enumList[curr], enumList[i]];
curr++;
}
}
// assert enumList[start:curr] >= enumList[curr:end]

if (curr === K) {
// assert enumList[start:K] >= enumList[K:end]
break;
} else if (curr > K) {
// enumList[curr: end] is useless
// curr == K <-> curr-1 is the kth,
// but the kth is at a smalller index,
// so curr-1 is (k+1)th or worse and can
// be thrown
end = curr - 1;
} else {
start = curr;
}
iters++;
}

return enumList.slice(0, K);
}
35 changes: 35 additions & 0 deletions lib/ops/topK.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT license.

import {Attribute} from '../attribute';
import {InferenceHandler} from '../backend';
import {FLOAT_TYPES, Operator} from '../operators';
import {Tensor} from '../tensor';

export abstract class TopK implements Operator {
abstract run(inferenceHandler: InferenceHandler, inputs: Tensor[]): Tensor[]|Promise<Tensor[]>;

initialize(attributes: Attribute): void {
this.axis = attributes.getInt('axis', -1);
this.k = attributes.getInt('k');
}

checkInputs(inputs: Tensor[]): boolean {
if (!inputs || inputs.length !== 1) {
return false;
}

return this.checkInputTypes(inputs);
}

protected checkInputTypes(inputs: Tensor[]): boolean {
if (FLOAT_TYPES.indexOf(inputs[0].type) === -1) {
return false;
}

return true;
}

protected axis: number;
protected k: number;
}
136 changes: 136 additions & 0 deletions test/data/ops/topk.jsonc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,136 @@
[
{
"name": "topk",
"operator": "TopK",
"attributes": [
{ "name": "axis", "data": 0, "type": "int" },
{ "name": "k", "data": 1, "type": "int" }
],
"cases": [
{
"name": "T[0] k=1",
"inputs": [
{
"data": [-3, 2],
"dims": [2],
"type": "float32"
}
],
"outputs": [
{
"data": [2],
"dims": [1],
"type": "float32"
},
{
"data": [1],
"dims": [1],
"type": "int32"
}
]
},
{
"name": "T[1] k=1",
"inputs": [
{
"data": [-3, 2, 1, -1],
"dims": [2, 2],
"type": "float32"
}
],
"outputs": [
{
"data": [1, 2],
"dims": [1, 2],
"type": "float32"
},
{
"data": [1, 0],
"dims": [1, 2],
"type": "int32"
}
]
},
{
"name": "T[2] k=1",
"inputs": [
{
"data": [3, 3, 3, 3, 3, 3],
"dims": [3, 2],
"type": "float32"
}
],
"outputs": [
{
"data": [3, 3],
"dims": [1, 2],
"type": "float32"
},
{
"data": [0, 0],
"dims": [1, 2],
"type": "int32"
}
]
}
]
},
{
"name": "topk",
"operator": "TopK",
"attributes": [
{ "name": "axis", "data": 1, "type": "int" },
{ "name": "k", "data": 2, "type": "int" }
],
"cases": [
{
"name": "T[0] k=2",
"inputs": [
{
"data": [
3,
3,
3,
3,
1,
1,
2,
3,
2,
3,
3,
3,
-1,
2,
3,
4,
5,
6,
5,
6,
3,
-2,
-3,
5
],
"dims": [2, 3, 4],
"type": "float32"
}
],
"outputs": [
{
"data": [3, 3, 3, 3, 2, 3, 3, 3, 5, 6, 5, 6, 3, 2, 3, 5],

"dims": [2, 2, 4],
"type": "float32"
},
{
"data": [0, 0, 0, 0, 2, 2, 2, 1, 1, 1, 1, 1, 2, 0, 0, 2],
"dims": [2, 2, 4],
"type": "int32"
}
]
}
]
}
]