1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
|
#pragma once
#include "caffe2/operators/spatial_batch_norm_op.h"
#include "caffe2/quantization/server/dnnlowp_op.h"
namespace caffe2 {
/**
* Note this implementation assumes SCALE, BIAS, EST_MEAN, and EST_VAR inputs
* are still in fp32, so is epsilon argument
*/
template <typename T, bool ReluFused = false>
class SpatialBNDNNLowPOp final : public DNNLowPOp<T, SpatialBNOp<CPUContext>> {
public:
USE_OPERATOR_FUNCTIONS(CPUContext);
USE_DNNLOWP_OPERATOR_BASE_FUNCTIONS(T, SpatialBNOp<CPUContext>);
SpatialBNDNNLowPOp(const OperatorDef& operator_def, Workspace* ws);
~SpatialBNDNNLowPOp() override = default;
bool RunOnDevice() override;
private:
void ComputeFusedParam_(
const int C,
const float* scale,
const float* bias,
const float* mean,
const float* var,
float* alpha,
float* beta);
double epsilon_;
const StorageOrder order_;
Tensor alpha_;
Tensor beta_;
INPUT_TAGS(INPUT, SCALE, BIAS, EST_MEAN, EST_VAR);
OUTPUT_TAGS(OUTPUT);
};
namespace internal {
template <typename T>
void SpatialBNNHWCAVX2(
const int N,
const int C,
const int HxW,
const int in_zero_point,
const int out_zero_point,
const T* X,
const float* alpha,
const float* beta,
T* Y,
bool relu_fused);
} // namespace internal
} // namespace caffe2
|
