| 26 | namespace TEngine { |
| 27 | |
| 28 | bool Permute::InferShape(const std::vector<TEngine::TShape>& ishape, std::vector<TEngine::TShape>& oshape, int layout) |
| 29 | { |
| 30 | const TShape& input = ishape[0]; |
| 31 | const std::vector<int> dims = input.GetDim(); |
| 32 | |
| 33 | // only support for 0231[bhwc] |
| 34 | if((param_.order0 == 0) && (param_.order1 == 2) && (param_.order2 == 3) && (param_.order3 == 1)) |
| 35 | { |
| 36 | int n = input.GetN(); |
| 37 | int c = input.GetC(); |
| 38 | int h = input.GetH(); |
| 39 | int w = input.GetW(); |
| 40 | TShape shape; |
| 41 | std::vector<int> dim = {n, h, w, c}; |
| 42 | shape.SetDim(dim); |
| 43 | shape.SetDataLayout(TENGINE_LAYOUT_NHWC); |
| 44 | oshape[0] = shape; |
| 45 | return true; |
| 46 | } |
| 47 | else if((param_.order0 == 1) && (param_.order1 == 0) && (param_.order2 == 2) && dims.size() == 3) |
| 48 | { |
| 49 | // int n = input.GetN(); |
| 50 | int c = input.Shape(0); |
| 51 | int h = input.Shape(1); |
| 52 | int w = input.Shape(2); |
| 53 | TShape shape; |
| 54 | std::vector<int> dim = {h, c, w}; |
| 55 | shape.SetDim(dim); |
| 56 | oshape[0] = shape; |
| 57 | return true; |
| 58 | } |
| 59 | else |
| 60 | { |
| 61 | return false; |
| 62 | } |
| 63 | } |
| 64 | |
| 65 | void Permute::SetSchema(void) |
| 66 | { |