Skip to content

GitLab

Commit e070ca45 authored by Xue Fuzhao's avatar Xue Fuzhao Committed by Frank Lee
Browse files

[NFC] polish colossalai/fx/tracer/meta_patch/patched_module/convolution.py code style (#1563)

parent 9823cbf2
Hide whitespace changes
Inline Side-by-side
Showing with 16 additions and 19 deletions
colossalai/fx/tracer/meta_patch/patched_module/convolution.py
View file @ e070ca45
...@@ -56,33 +56,32 @@ def torch_nn_conv3d(self, input): ...@@ -56,33 +56,32 @@ def torch_nn_conv3d(self, input):
) )
return torch.empty(result_shape, device='meta') return torch.empty(result_shape, device='meta')
@meta_patched_module.register(torch.nn.ConvTranspose1d) @meta_patched_module.register(torch.nn.ConvTranspose1d)
def torch_nn_convtranspose1d(self, input): def torch_nn_convtranspose1d(self, input):
# the output shape is calculated using the formula stated # the output shape is calculated using the formula stated
# at https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose1d.html # at https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose1d.html
l_in = input.shape[-1] l_in = input.shape[-1]
c_out = self.out_channels c_out = self.out_channels
l_out = math.floor((l_in - 1) * self.stride[0] - 2 * self.padding[0] + l_out = math.floor((l_in - 1) * self.stride[0] - 2 * self.padding[0] + self.dilation[0] *
self.dilation[0] * (self.kernel_size[0] - 1) + (self.kernel_size[0] - 1) + self.output_padding[0] + 1)
self.output_padding[0] + 1)
result_shape = input.shape[:-2] + ( result_shape = input.shape[:-2] + (
c_out, c_out,
l_out, l_out,
) )
return torch.empty(result_shape, device='meta') return torch.empty(result_shape, device='meta')
@meta_patched_module.register(torch.nn.ConvTranspose2d) @meta_patched_module.register(torch.nn.ConvTranspose2d)
def torch_nn_convtranspose2d(self, input): def torch_nn_convtranspose2d(self, input):
# the output shape is calculated using the formula stated # the output shape is calculated using the formula stated
# at https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose2d.html # at https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose2d.html
h_in, w_in = input.shape[-2:] h_in, w_in = input.shape[-2:]
c_out = self.out_channels c_out = self.out_channels
h_out = math.floor((h_in - 1) * self.stride[0] - 2 * self.padding[0] + h_out = math.floor((h_in - 1) * self.stride[0] - 2 * self.padding[0] + self.dilation[0] *
self.dilation[0] * (self.kernel_size[0] - 1) + (self.kernel_size[0] - 1) + self.output_padding[0] + 1)
self.output_padding[0] + 1) w_out = math.floor((w_in - 1) * self.stride[1] - 2 * self.padding[1] + self.dilation[1] *
w_out = math.floor((w_in - 1) * self.stride[1] - 2 * self.padding[1] + (self.kernel_size[1] - 1) + self.output_padding[1] + 1)
self.dilation[1] * (self.kernel_size[1] - 1) +
self.output_padding[1] + 1)
result_shape = input.shape[:-3] + ( result_shape = input.shape[:-3] + (
c_out, c_out,
h_out, h_out,
...@@ -90,25 +89,23 @@ def torch_nn_convtranspose2d(self, input): ...@@ -90,25 +89,23 @@ def torch_nn_convtranspose2d(self, input):
) )
return torch.empty(result_shape, device='meta') return torch.empty(result_shape, device='meta')
@meta_patched_module.register(torch.nn.ConvTranspose3d) @meta_patched_module.register(torch.nn.ConvTranspose3d)
def torch_nn_convtranspose3d(self, input): def torch_nn_convtranspose3d(self, input):
# the output shape is calculated using the formula stated # the output shape is calculated using the formula stated
# at https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose3d.html # at https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose3d.html
d_in, h_in, w_in = input.shape[-3:] d_in, h_in, w_in = input.shape[-3:]
c_out = self.out_channels c_out = self.out_channels
d_out = math.floor((d_in - 1) * self.stride[0] - 2 * self.padding[0] + d_out = math.floor((d_in - 1) * self.stride[0] - 2 * self.padding[0] + self.dilation[0] *
self.dilation[0] * (self.kernel_size[0] - 1) + (self.kernel_size[0] - 1) + self.output_padding[0] + 1)
self.output_padding[0] + 1) h_out = math.floor((h_in - 1) * self.stride[1] - 2 * self.padding[1] + self.dilation[1] *
h_out = math.floor((h_in - 1) * self.stride[1] - 2 * self.padding[1] + (self.kernel_size[1] - 1) + self.output_padding[1] + 1)
self.dilation[1] * (self.kernel_size[1] - 1) + w_out = math.floor((w_in - 1) * self.stride[2] - 2 * self.padding[2] + self.dilation[2] *
self.output_padding[1] + 1) (self.kernel_size[2] - 1) + self.output_padding[2] + 1)
w_out = math.floor((w_in - 1) * self.stride[2] - 2 * self.padding[2] +
self.dilation[2] * (self.kernel_size[2] - 1) +
self.output_padding[2] + 1)
result_shape = input.shape[:-4] + ( result_shape = input.shape[:-4] + (
c_out, c_out,
d_out, d_out,
h_out, h_out,
w_out, w_out,
) )
return torch.empty(result_shape, device='meta') return torch.empty(result_shape, device='meta')
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment