Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Conv padding: remove string values, move to TF-style values #62

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 30 additions & 28 deletions nir/ir.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ def _parse_shape_argument(x: Types, key: str):

def _calculate_conv_output(
input_shape: typing.Union[int, typing.Sequence[int]],
padding: typing.Union[int, str, typing.Sequence[int]],
padding: np.ndarray,
dilation: typing.Union[int, typing.Sequence[int]],
kernel_size: typing.Union[int, typing.Sequence[int]],
stride: typing.Union[int, typing.Sequence[int]],
Expand All @@ -37,8 +37,8 @@ def _calculate_conv_output(

:param input_shape: input shape, either int or (int, int)
:type input_shape: int | typing.Sequence[int]
:param padding: padding
:type padding: int | typing.Sequence[int]
:param padding: padding, each dimension has two values
:type padding: integer array, of shape (2,) or (2, 2), etc.
:param dilation: dilation
:type dilation: int | typing.Sequence[int]
:param kernel_size: kernel size
Expand All @@ -53,23 +53,18 @@ def _calculate_conv_output(
ndim = 1
else:
ndim = len(input_shape)
if isinstance(padding, str) and padding == 'valid':
padding = [0] * ndim
shapes = []
for i in range(ndim):
if isinstance(padding, str) and padding == 'same':
shape = input_shape[i]
else:
shape = np.floor(
(
_index_tuple(input_shape, i)
+ 2 * _index_tuple(padding, i)
- _index_tuple(dilation, i) * (_index_tuple(kernel_size, i) - 1)
- 1
)
/ _index_tuple(stride, i)
+ 1
shape = np.floor(
(
_index_tuple(input_shape, i)
+ sum(padding if isinstance(padding[0], (int, np.integer)) else padding[i])
- _index_tuple(dilation, i) * (_index_tuple(kernel_size, i) - 1)
- 1
)
/ _index_tuple(stride, i)
+ 1
)
shapes.append(int(shape))
return np.array(shapes)

Expand Down Expand Up @@ -352,8 +347,8 @@ class Conv1d(NIRNode):
:type weight: np.ndarray
:param stride: Stride
:type stride: int
:param padding: Padding, if string must be 'same' or 'valid'
:type padding: int | str
:param padding: Padding
:type padding: int
:param dilation: Dilation
:type dilation: int
:param groups: Groups
Expand All @@ -365,14 +360,17 @@ class Conv1d(NIRNode):
input_shape: typing.Optional[int] # N
weight: np.ndarray # Weight C_out * C_in * N
stride: int # Stride
padding: typing.Union[int, str] # Padding
padding: typing.Union[int, typing.Tuple[int, int]] # Padding
dilation: int # Dilation
groups: int # Groups
bias: np.ndarray # Bias C_out

def __post_init__(self):
if isinstance(self.padding, str) and self.padding not in ["same", "valid"]:
raise ValueError(f"padding must be 'same', 'valid', or int, not {self.padding}")
if isinstance(self.padding, int):
self.padding = (self.padding, self.padding)
elif isinstance(self.padding[0], (int, np.integer)):
assert len(self.padding) == 2, "Conv1D Padding must be of length 2"
self.padding = tuple([int(p) for p in self.padding])
if self.input_shape is None:
# leave input and output types undefined
self.input_type = {"input": None}
Expand Down Expand Up @@ -407,8 +405,8 @@ class Conv2d(NIRNode):
:type weight: np.ndarray
:param stride: Stride
:type stride: int | int, int
:param padding: Padding, if string must be 'same' or 'valid'
:type padding: int | int, int | str
:param padding: Padding
:type padding: int | int, int | ((int, int), (int, int))
:param dilation: Dilation
:type dilation: int | int, int
:param groups: Groups
Expand All @@ -421,16 +419,20 @@ class Conv2d(NIRNode):
input_shape: typing.Optional[typing.Tuple[int, int]] # N_x, N_y
weight: np.ndarray # Weight C_out * C_in * W_x * W_y
stride: typing.Union[int, typing.Tuple[int, int]] # Stride
padding: typing.Union[int, typing.Tuple[int, int], str] # Padding
padding: typing.Union[int, typing.Tuple[int, int], typing.Sequence[typing.Tuple[int, int]]]
dilation: typing.Union[int, typing.Tuple[int, int]] # Dilation
groups: int # Groups
bias: np.ndarray # Bias C_out

def __post_init__(self):
if isinstance(self.padding, str) and self.padding not in ["same", "valid"]:
raise ValueError(f"padding must be 'same', 'valid', or int, not {self.padding}")
if isinstance(self.padding, int):
self.padding = (self.padding, self.padding)
self.padding = ((self.padding, self.padding), (self.padding, self.padding))
elif isinstance(self.padding[0], (int, np.integer)):
assert len(self.padding) == 2, "Conv2D Padding must be of length 2"
self.padding = (self.padding[0], self.padding[0]), (self.padding[1], self.padding[1])
elif isinstance(self.padding[0][0], (int, np.integer)):
assert all(len(p) == 2 for p in self.padding), "Padding must be a tuple of length 2"
self.padding = tuple([tuple([int(e) for e in p]) for p in self.padding])
if isinstance(self.stride, int):
self.stride = (self.stride, self.stride)
if isinstance(self.dilation, int):
Expand Down