| 252 | typedef Eigen::ThreadPoolDevice CPUDevice; |
| 253 | |
| 254 | class QuantizedInstanceNorm : public OpKernel { |
| 255 | public: |
| 256 | explicit QuantizedInstanceNorm(OpKernelConstruction* context) |
| 257 | : OpKernel(context) { |
| 258 | OP_REQUIRES_OK(context, |
| 259 | context->GetAttr("variance_epsilon", &variance_epsilon_)); |
| 260 | OP_REQUIRES_OK(context, |
| 261 | context->GetAttr("min_separation", &min_separation_)); |
| 262 | OP_REQUIRES_OK( |
| 263 | context, context->GetAttr("output_range_given", &output_range_given_)); |
| 264 | if (output_range_given_) { |
| 265 | OP_REQUIRES_OK(context, context->GetAttr("given_y_min", &given_y_min_)); |
| 266 | OP_REQUIRES_OK(context, context->GetAttr("given_y_max", &given_y_max_)); |
| 267 | OP_REQUIRES(context, given_y_min_ < given_y_max_, |
| 268 | errors::InvalidArgument( |
| 269 | "given_y_min must be less than given_y_max : ", |
| 270 | given_y_min_, " >= ", given_y_max_)); |
| 271 | } |
| 272 | } |
| 273 | |
| 274 | void Compute(OpKernelContext* context) override { |
| 275 | const Tensor& input = context->input(0); |
| 276 | |
| 277 | float input_min = context->input(1).flat<float>()(0); |
| 278 | float input_max = context->input(2).flat<float>()(0); |
| 279 | float input_scale = (input_max - input_min) / 255.0f; |
| 280 | |
| 281 | OP_REQUIRES(context, input_min < input_max, |
| 282 | errors::InvalidArgument( |
| 283 | "input_min must be less than input_max : ", input_min, |
| 284 | " >= ", input_max)); |
| 285 | |
| 286 | auto input_tensor = input.tensor<quint8, 4>(); |
| 287 | auto N = input_tensor.dimension(0); |
| 288 | auto H = input_tensor.dimension(1); |
| 289 | auto W = input_tensor.dimension(2); |
| 290 | auto C = input_tensor.dimension(3); |
| 291 | |
| 292 | Tensor* output = nullptr; |
| 293 | OP_REQUIRES_OK(context, |
| 294 | context->allocate_output(0, input.shape(), &output)); |
| 295 | |
| 296 | Tensor* output_min = nullptr; |
| 297 | OP_REQUIRES_OK(context, context->allocate_output(1, {}, &output_min)); |
| 298 | Tensor* output_max = nullptr; |
| 299 | OP_REQUIRES_OK(context, context->allocate_output(2, {}, &output_max)); |
| 300 | |
| 301 | typedef TTypes<float>::Tensor::Index Index; |
| 302 | |
| 303 | #if defined(EIGEN_HAS_INDEX_LIST) |
| 304 | const Eigen::IndexList<Eigen::type2index<1>, Eigen::type2index<2>> |
| 305 | reduction_indices; |
| 306 | Eigen::IndexList<Eigen::type2index<1>, Index, Index, Eigen::type2index<1>> |
| 307 | broadcast_spec; |
| 308 | broadcast_spec.set(1, H); |
| 309 | broadcast_spec.set(2, W); |
| 310 | Eigen::IndexList<Index, Eigen::type2index<1>, Eigen::type2index<1>, Index> |
| 311 | expand_spec; |