| 61 | } |
| 62 | |
| 63 | SmallVector<TensorPtr> apply_on_physical_tensor( |
| 64 | const OpDef& def, const SmallVector<TensorPtr>& inputs, |
| 65 | SmallVector<LogicalTensorDesc>& output_descs, const bool& validated) { |
| 66 | auto&& op_def = def.cast_final_safe<InstanceNorm>(); |
| 67 | size_t nr_inp = inputs.size(); |
| 68 | auto p = op_def.param(); |
| 69 | |
| 70 | mgb_assert( |
| 71 | (nr_inp == 3 && p.affine) || (nr_inp == 1 && !p.affine), |
| 72 | "num of inputs of instancenorm should be 1 or 3 but you give %zu", |
| 73 | inputs.size()); |
| 74 | |
| 75 | auto cn = inputs[0]->comp_node(); |
| 76 | using Format = megdnn::param::GroupNorm::Format; |
| 77 | mgb_assert(p.format == Format::NCHW, "only support inputs in shape NCHW."); |
| 78 | size_t C = inputs[0]->shape()[1]; |
| 79 | p.group = C; |
| 80 | |
| 81 | DnnOprCaller<megdnn::GroupNorm> caller(cn, p); |
| 82 | |
| 83 | auto&& [oup_layout, mean_layout, rstd_layout] = caller.deduce_layouts<3>( |
| 84 | inputs[0]->layout(), TensorLayout{}, TensorLayout{}); |
| 85 | |
| 86 | auto out = Tensor::make(oup_layout, cn); |
| 87 | auto mean = Tensor::make(mean_layout, cn); |
| 88 | auto rstd = Tensor::make(rstd_layout, cn); |
| 89 | |
| 90 | if (oup_layout.is_empty()) { |
| 91 | return {out, mean, rstd}; |
| 92 | } |
| 93 | |
| 94 | if (p.affine) { |
| 95 | caller.exec_with_ws(inputs[0], inputs[1], inputs[2], out, mean, rstd); |
| 96 | } else { |
| 97 | megdnn::TensorND empty_dnn; |
| 98 | caller.exec_with_ws(inputs[0], empty_dnn, empty_dnn, out, mean, rstd); |
| 99 | } |
| 100 | return {out, mean, rstd}; |
| 101 | } |
| 102 | |
| 103 | OP_TRAIT_REG(InstanceNorm, InstanceNorm) |
| 104 | .apply_on_var_node(apply_on_var_node) |