onnxruntime

1
// Copyright (c) Microsoft Corporation. All rights reserved.
2
// Licensed under the MIT License.
3

4
import { Attribute } from '../../../attribute';
5
import { MAX_CLIP, MIN_CLIP } from '../../../util';
6
import { GlslValueFunction } from '../glsl-definitions';
7

8
import { glslClip, glslRelu, glslSigmoid } from './unary-op';
9

10
export interface InternalActivationAttributes {
11
  readonly activation: string;
12
  readonly clipMin?: number;
13
  readonly clipMax?: number;
14
  readonly activationCacheKey: string;
15
}
16

17
export function getActivationSnippet(attributes: InternalActivationAttributes) {
18
  let func: GlslValueFunction;
19
  switch (attributes.activation) {
20
    case 'Relu':
21
      func = glslRelu();
22
      break;
23
    case 'Sigmoid':
24
      func = glslSigmoid();
25
      break;
26
    case 'Clip':
27
      func = glslClip(attributes.clipMin!, attributes.clipMax!);
28
      break;
29
    // TODO: adding other activations that can be fused.
30
    default:
31
      return { activationFunction: '', applyActivation: '' };
32
  }
33

34
  const activationName = func.name;
35
  const activationFunction = func.body;
36
  const applyActivation = `value = ${activationName}_(value);`;
37
  return { activationFunction, applyActivation };
38
}
39

40
export const parseInternalActivationAttributes = (attributes: Attribute): InternalActivationAttributes => {
41
  const activation = attributes.getString('activation', '');
42

43
  if (activation === 'Clip') {
44
    const [clipMin, clipMax] = attributes.getFloats('activation_params', [MIN_CLIP, MAX_CLIP]);
45
    return { activation, clipMax, clipMin, activationCacheKey: `${activation}:${clipMin},${clipMax}` };
46
  }
47
  return { activation, activationCacheKey: activation };
48
};
49