patrik.gustavsson | eeb8515 | 2020-12-21 17:10:40 +0000 | [diff] [blame] | 1 | # Copyright (C) 2020 Arm Limited or its affiliates. All rights reserved. |
| 2 | # |
| 3 | # SPDX-License-Identifier: Apache-2.0 |
| 4 | # |
| 5 | # Licensed under the Apache License, Version 2.0 (the License); you may |
| 6 | # not use this file except in compliance with the License. |
| 7 | # You may obtain a copy of the License at |
| 8 | # |
| 9 | # www.apache.org/licenses/LICENSE-2.0 |
| 10 | # |
| 11 | # Unless required by applicable law or agreed to in writing, software |
| 12 | # distributed under the License is distributed on an AS IS BASIS, WITHOUT |
| 13 | # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 14 | # See the License for the specific language governing permissions and |
| 15 | # limitations under the License. |
| 16 | # Description: |
| 17 | # Defines the class Shape4D. |
Tim Hall | 73e843f | 2021-02-04 22:47:46 +0000 | [diff] [blame] | 18 | from collections import namedtuple |
| 19 | |
patrik.gustavsson | eeb8515 | 2020-12-21 17:10:40 +0000 | [diff] [blame] | 20 | from .numeric_util import full_shape |
Tim Hall | 73e843f | 2021-02-04 22:47:46 +0000 | [diff] [blame] | 21 | from .numeric_util import round_up_divide |
patrik.gustavsson | eeb8515 | 2020-12-21 17:10:40 +0000 | [diff] [blame] | 22 | |
| 23 | |
Tim Hall | 73e843f | 2021-02-04 22:47:46 +0000 | [diff] [blame] | 24 | class Shape4D(namedtuple("Shape4D", ["batch", "height", "width", "depth"])): |
patrik.gustavsson | eeb8515 | 2020-12-21 17:10:40 +0000 | [diff] [blame] | 25 | """ |
| 26 | 4D Shape (in NHWC format) |
| 27 | """ |
| 28 | |
Tim Hall | 73e843f | 2021-02-04 22:47:46 +0000 | [diff] [blame] | 29 | def __new__(cls, n=1, h=1, w=1, c=1): |
| 30 | assert n is not None |
| 31 | if isinstance(n, list): |
| 32 | assert h == 1 and w == 1 and c == 1 |
| 33 | tmp = full_shape(4, n, 1) |
| 34 | self = super(Shape4D, cls).__new__(cls, tmp[0], tmp[1], tmp[2], tmp[3]) |
| 35 | else: |
| 36 | self = super(Shape4D, cls).__new__(cls, n, h, w, c) |
| 37 | return self |
| 38 | |
| 39 | @classmethod |
| 40 | def from_list(cls, shape, base=1): |
| 41 | tmp = full_shape(4, shape, base) |
| 42 | return cls(tmp[0], tmp[1], tmp[2], tmp[3]) |
| 43 | |
| 44 | @classmethod |
| 45 | def from_hwc(cls, h, w, c): |
| 46 | return cls(1, h, w, c) |
| 47 | |
| 48 | def with_batch(self, new_batch): |
| 49 | return Shape4D(new_batch, self.height, self.width, self.depth) |
| 50 | |
| 51 | def with_height(self, new_height): |
| 52 | return Shape4D(self.batch, new_height, self.width, self.depth) |
| 53 | |
| 54 | def with_width(self, new_width): |
| 55 | return Shape4D(self.batch, self.height, new_width, self.depth) |
| 56 | |
| 57 | def with_hw(self, new_height, new_width): |
| 58 | return Shape4D(self.batch, new_height, new_width, self.depth) |
| 59 | |
| 60 | def with_depth(self, new_depth): |
| 61 | return Shape4D(self.batch, self.height, self.width, new_depth) |
| 62 | |
| 63 | def add(self, n, h, w, c): |
| 64 | return Shape4D(self.batch + n, self.height + h, self.width + w, self.depth + c) |
| 65 | |
| 66 | def __add__(self, rhs): |
| 67 | return Shape4D(self.batch + rhs.batch, self.height + rhs.height, self.width + rhs.width, self.depth + rhs.depth) |
| 68 | |
| 69 | def __sub__(self, rhs): |
| 70 | return Shape4D(self.batch - rhs.batch, self.height - rhs.height, self.width - rhs.width, self.depth - rhs.depth) |
| 71 | |
| 72 | def __floordiv__(self, rhs): |
| 73 | return Shape4D( |
| 74 | self.batch // rhs.batch, self.height // rhs.height, self.width // rhs.width, self.depth // rhs.depth |
| 75 | ) |
| 76 | |
| 77 | def __mod__(self, rhs): |
| 78 | return Shape4D(self.batch % rhs.batch, self.height % rhs.height, self.width % rhs.width, self.depth % rhs.depth) |
patrik.gustavsson | eeb8515 | 2020-12-21 17:10:40 +0000 | [diff] [blame] | 79 | |
| 80 | def __str__(self): |
Tim Hall | 73e843f | 2021-02-04 22:47:46 +0000 | [diff] [blame] | 81 | return f"<Shape4D {list(self)}>" |
patrik.gustavsson | eeb8515 | 2020-12-21 17:10:40 +0000 | [diff] [blame] | 82 | |
Tim Hall | 73e843f | 2021-02-04 22:47:46 +0000 | [diff] [blame] | 83 | def div_round_up(self, rhs): |
| 84 | return Shape4D( |
| 85 | round_up_divide(self.batch, rhs.batch), |
| 86 | round_up_divide(self.height, rhs.height), |
| 87 | round_up_divide(self.width, rhs.width), |
| 88 | round_up_divide(self.depth, rhs.depth), |
| 89 | ) |
patrik.gustavsson | eeb8515 | 2020-12-21 17:10:40 +0000 | [diff] [blame] | 90 | |
Tim Hall | 73e843f | 2021-02-04 22:47:46 +0000 | [diff] [blame] | 91 | def elements(self): |
| 92 | return self.batch * self.width * self.height * self.depth |
patrik.gustavsson | eeb8515 | 2020-12-21 17:10:40 +0000 | [diff] [blame] | 93 | |
Tim Hall | 73e843f | 2021-02-04 22:47:46 +0000 | [diff] [blame] | 94 | def elements_wh(self): |
| 95 | return self.width * self.height |
patrik.gustavsson | eeb8515 | 2020-12-21 17:10:40 +0000 | [diff] [blame] | 96 | |
Tim Hall | 73e843f | 2021-02-04 22:47:46 +0000 | [diff] [blame] | 97 | def is_empty(self): |
| 98 | return (self.batch + self.width + self.height + self.depth) == 0 |
patrik.gustavsson | eeb8515 | 2020-12-21 17:10:40 +0000 | [diff] [blame] | 99 | |
| 100 | def as_list(self): |
Tim Hall | 73e843f | 2021-02-04 22:47:46 +0000 | [diff] [blame] | 101 | return list(self) |
Patrik Gustavsson | 3a26920 | 2021-01-21 08:28:55 +0100 | [diff] [blame] | 102 | |
| 103 | def get_hw_as_list(self): |
| 104 | return list([self.height, self.width]) |