MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / apply_on_physical_tensor

Function apply_on_physical_tensor

imperative/src/impl/ops/instance_norm.cpp:63–101  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

61}
62
63SmallVector<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
103OP_TRAIT_REG(InstanceNorm, InstanceNorm)
104 .apply_on_var_node(apply_on_var_node)

Callers

nothing calls this directly

Calls 8

exec_with_wsMethod · 0.80
makeFunction · 0.70
sizeMethod · 0.45
paramMethod · 0.45
comp_nodeMethod · 0.45
shapeMethod · 0.45
layoutMethod · 0.45
is_emptyMethod · 0.45

Tested by

no test coverage detected