mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-03 10:06:44 +00:00
[misc] update pre-commit and run all files (#4752)
* [misc] update pre-commit * [misc] run pre-commit * [misc] remove useless configuration files * [misc] ignore cuda for clang-format
This commit is contained in:
Hongxin Liu
GitHub
@@ -36,12 +36,12 @@ def _assert_output_shape(data, module, patch_fn, expect_exception, output_shape)
|
||||
@clear_cache_before_run()
|
||||
def test_linear():
|
||||
# test linear patch can produce the meta output with correct shape
|
||||
data = torch.rand(2, 4, device='meta')
|
||||
data = torch.rand(2, 4, device="meta")
|
||||
module = torch.nn.Linear(4, 2)
|
||||
_assert_output_shape(data, module, patched_module.torch_nn_linear, False, torch.Size([2, 2]))
|
||||
|
||||
# test if the linear patch can catch exception when dimension does not match
|
||||
data = torch.rand(2, 2, device='meta')
|
||||
data = torch.rand(2, 2, device="meta")
|
||||
_assert_output_shape(data, module, patched_module.torch_nn_linear, True, None)
|
||||
|
||||
|
||||
@@ -51,20 +51,20 @@ def test_rnn():
|
||||
data = (torch.randn(5, 3, 10), torch.randn(2, 3, 20))
|
||||
module = torch.nn.RNN(10, 20, 2)
|
||||
output, hn = module(*data)
|
||||
meta_data = (torch.randn(5, 3, 10).to('meta'), torch.randn(2, 3, 20).to('meta'))
|
||||
meta_data = (torch.randn(5, 3, 10).to("meta"), torch.randn(2, 3, 20).to("meta"))
|
||||
_assert_output_shape(meta_data, module, patched_module.torch_nn_rnn, False, (output.shape, hn.shape))
|
||||
|
||||
# test if the rnn patch can catch exception when dimension does not match
|
||||
data = (torch.randn(5, 3, 10), torch.randn(2, 3, 20))
|
||||
module = torch.nn.RNN(10, 20, 2)
|
||||
output, hn = module(*data)
|
||||
meta_data = (torch.randn(5, 3, 1).to('meta'), torch.randn(2, 3, 20).to('meta'))
|
||||
meta_data = (torch.randn(5, 3, 1).to("meta"), torch.randn(2, 3, 20).to("meta"))
|
||||
_assert_output_shape(meta_data, module, patched_module.torch_nn_rnn, True, None)
|
||||
|
||||
|
||||
@clear_cache_before_run()
|
||||
def test_embedding():
|
||||
data = torch.rand(2, 4, device='meta')
|
||||
data = torch.rand(2, 4, device="meta")
|
||||
|
||||
# test layernorm
|
||||
ln = torch.nn.LayerNorm(4)
|
||||
@@ -76,67 +76,71 @@ def test_embedding():
|
||||
|
||||
# test batch norm 1d
|
||||
bn1d = torch.nn.BatchNorm1d(4)
|
||||
data = torch.rand(2, 4, device='meta')
|
||||
_assert_output_shape(data=data,
|
||||
module=bn1d,
|
||||
patch_fn=patched_module.torch_nn_normalize,
|
||||
expect_exception=False,
|
||||
output_shape=data.shape)
|
||||
data = torch.rand(2, 4, device="meta")
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=bn1d,
|
||||
patch_fn=patched_module.torch_nn_normalize,
|
||||
expect_exception=False,
|
||||
output_shape=data.shape,
|
||||
)
|
||||
|
||||
data = torch.rand(2, 4, device='meta')
|
||||
_assert_output_shape(data=data,
|
||||
module=bn1d,
|
||||
patch_fn=patched_module.torch_nn_normalize,
|
||||
expect_exception=False,
|
||||
output_shape=data.shape)
|
||||
data = torch.rand(2, 4, device="meta")
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=bn1d,
|
||||
patch_fn=patched_module.torch_nn_normalize,
|
||||
expect_exception=False,
|
||||
output_shape=data.shape,
|
||||
)
|
||||
|
||||
data = torch.rand(2, 3, 4, device='meta')
|
||||
_assert_output_shape(data=data,
|
||||
module=bn1d,
|
||||
patch_fn=patched_module.torch_nn_normalize,
|
||||
expect_exception=False,
|
||||
output_shape=data.shape)
|
||||
data = torch.rand(2, 3, 4, device="meta")
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=bn1d,
|
||||
patch_fn=patched_module.torch_nn_normalize,
|
||||
expect_exception=False,
|
||||
output_shape=data.shape,
|
||||
)
|
||||
|
||||
data = torch.rand(1, 2, 3, 4, device='meta')
|
||||
_assert_output_shape(data=data,
|
||||
module=bn1d,
|
||||
patch_fn=patched_module.torch_nn_normalize,
|
||||
expect_exception=True,
|
||||
output_shape=None)
|
||||
data = torch.rand(1, 2, 3, 4, device="meta")
|
||||
_assert_output_shape(
|
||||
data=data, module=bn1d, patch_fn=patched_module.torch_nn_normalize, expect_exception=True, output_shape=None
|
||||
)
|
||||
|
||||
# test batch norm 2d
|
||||
bn2d = torch.nn.BatchNorm2d(4)
|
||||
|
||||
data = torch.rand(1, 2, 3, 4, device='meta')
|
||||
_assert_output_shape(data=data,
|
||||
module=bn2d,
|
||||
patch_fn=patched_module.torch_nn_normalize,
|
||||
expect_exception=False,
|
||||
output_shape=data.shape)
|
||||
data = torch.rand(1, 2, 3, 4, device="meta")
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=bn2d,
|
||||
patch_fn=patched_module.torch_nn_normalize,
|
||||
expect_exception=False,
|
||||
output_shape=data.shape,
|
||||
)
|
||||
|
||||
data = torch.rand(2, 3, 4, device='meta')
|
||||
_assert_output_shape(data=data,
|
||||
module=bn2d,
|
||||
patch_fn=patched_module.torch_nn_normalize,
|
||||
expect_exception=True,
|
||||
output_shape=None)
|
||||
data = torch.rand(2, 3, 4, device="meta")
|
||||
_assert_output_shape(
|
||||
data=data, module=bn2d, patch_fn=patched_module.torch_nn_normalize, expect_exception=True, output_shape=None
|
||||
)
|
||||
|
||||
# # test batch size 3d
|
||||
bn3d = torch.nn.BatchNorm3d(4)
|
||||
|
||||
data = torch.rand(1, 1, 2, 3, 4, device='meta')
|
||||
_assert_output_shape(data=data,
|
||||
module=bn3d,
|
||||
patch_fn=patched_module.torch_nn_normalize,
|
||||
expect_exception=False,
|
||||
output_shape=data.shape)
|
||||
data = torch.rand(1, 1, 2, 3, 4, device="meta")
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=bn3d,
|
||||
patch_fn=patched_module.torch_nn_normalize,
|
||||
expect_exception=False,
|
||||
output_shape=data.shape,
|
||||
)
|
||||
|
||||
data = torch.rand(1, 2, 3, 4, device='meta')
|
||||
_assert_output_shape(data=data,
|
||||
module=bn3d,
|
||||
patch_fn=patched_module.torch_nn_normalize,
|
||||
expect_exception=True,
|
||||
output_shape=None)
|
||||
data = torch.rand(1, 2, 3, 4, device="meta")
|
||||
_assert_output_shape(
|
||||
data=data, module=bn3d, patch_fn=patched_module.torch_nn_normalize, expect_exception=True, output_shape=None
|
||||
)
|
||||
|
||||
|
||||
@clear_cache_before_run()
|
||||
@@ -146,35 +150,38 @@ def test_conv1d():
|
||||
|
||||
conv1d = torch.nn.Conv1d(in_channels=3, out_channels=4, kernel_size=2)
|
||||
materialized_output = conv1d(data)
|
||||
meta_data = data.to('meta')
|
||||
_assert_output_shape(data=meta_data,
|
||||
module=conv1d,
|
||||
patch_fn=patched_module.torch_nn_conv1d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
meta_data = data.to("meta")
|
||||
_assert_output_shape(
|
||||
data=meta_data,
|
||||
module=conv1d,
|
||||
patch_fn=patched_module.torch_nn_conv1d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
conv1d = torch.nn.Conv1d(in_channels=3, out_channels=4, kernel_size=2, padding=1)
|
||||
materialized_output = conv1d(data)
|
||||
meta_data = data.to('meta')
|
||||
_assert_output_shape(data=meta_data,
|
||||
module=conv1d,
|
||||
patch_fn=patched_module.torch_nn_conv1d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
meta_data = data.to("meta")
|
||||
_assert_output_shape(
|
||||
data=meta_data,
|
||||
module=conv1d,
|
||||
patch_fn=patched_module.torch_nn_conv1d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
conv1d = torch.nn.Conv1d(in_channels=3,
|
||||
out_channels=4,
|
||||
kernel_size=2,
|
||||
padding=1,
|
||||
dilation=2,
|
||||
padding_mode='reflect')
|
||||
conv1d = torch.nn.Conv1d(
|
||||
in_channels=3, out_channels=4, kernel_size=2, padding=1, dilation=2, padding_mode="reflect"
|
||||
)
|
||||
materialized_output = conv1d(data)
|
||||
meta_data = data.to('meta')
|
||||
_assert_output_shape(data=meta_data,
|
||||
module=conv1d,
|
||||
patch_fn=patched_module.torch_nn_conv1d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
meta_data = data.to("meta")
|
||||
_assert_output_shape(
|
||||
data=meta_data,
|
||||
module=conv1d,
|
||||
patch_fn=patched_module.torch_nn_conv1d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
|
||||
def test_conv2d():
|
||||
@@ -182,40 +189,45 @@ def test_conv2d():
|
||||
data = torch.rand(2, 3, 4, 4)
|
||||
conv2d = torch.nn.Conv2d(in_channels=3, out_channels=4, kernel_size=2)
|
||||
materialized_output = conv2d(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=conv2d,
|
||||
patch_fn=patched_module.torch_nn_conv2d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=conv2d,
|
||||
patch_fn=patched_module.torch_nn_conv2d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
conv2d = torch.nn.Conv2d(in_channels=3, out_channels=4, kernel_size=2, padding=1)
|
||||
materialized_output = conv2d(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=conv2d,
|
||||
patch_fn=patched_module.torch_nn_conv2d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=conv2d,
|
||||
patch_fn=patched_module.torch_nn_conv2d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
conv2d = torch.nn.Conv2d(in_channels=3, out_channels=4, kernel_size=2, padding=1, dilation=2)
|
||||
materialized_output = conv2d(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=conv2d,
|
||||
patch_fn=patched_module.torch_nn_conv2d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=conv2d,
|
||||
patch_fn=patched_module.torch_nn_conv2d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
conv2d = torch.nn.Conv2d(in_channels=3,
|
||||
out_channels=4,
|
||||
kernel_size=2,
|
||||
padding=1,
|
||||
dilation=2,
|
||||
padding_mode='reflect')
|
||||
conv2d = torch.nn.Conv2d(
|
||||
in_channels=3, out_channels=4, kernel_size=2, padding=1, dilation=2, padding_mode="reflect"
|
||||
)
|
||||
materialized_output = conv2d(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=conv2d,
|
||||
patch_fn=patched_module.torch_nn_conv2d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=conv2d,
|
||||
patch_fn=patched_module.torch_nn_conv2d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
|
||||
@clear_cache_before_run()
|
||||
@@ -224,40 +236,45 @@ def test_conv3d():
|
||||
data = torch.rand(2, 3, 4, 4, 4)
|
||||
conv3d = torch.nn.Conv3d(in_channels=3, out_channels=4, kernel_size=2)
|
||||
materialized_output = conv3d(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=conv3d,
|
||||
patch_fn=patched_module.torch_nn_conv3d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=conv3d,
|
||||
patch_fn=patched_module.torch_nn_conv3d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
conv3d = torch.nn.Conv3d(in_channels=3, out_channels=4, kernel_size=2, padding=1)
|
||||
materialized_output = conv3d(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=conv3d,
|
||||
patch_fn=patched_module.torch_nn_conv3d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=conv3d,
|
||||
patch_fn=patched_module.torch_nn_conv3d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
conv3d = torch.nn.Conv3d(in_channels=3, out_channels=4, kernel_size=2, padding=1, dilation=2)
|
||||
materialized_output = conv3d(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=conv3d,
|
||||
patch_fn=patched_module.torch_nn_conv3d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=conv3d,
|
||||
patch_fn=patched_module.torch_nn_conv3d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
conv3d = torch.nn.Conv3d(in_channels=3,
|
||||
out_channels=4,
|
||||
kernel_size=2,
|
||||
padding=1,
|
||||
dilation=2,
|
||||
padding_mode='reflect')
|
||||
conv3d = torch.nn.Conv3d(
|
||||
in_channels=3, out_channels=4, kernel_size=2, padding=1, dilation=2, padding_mode="reflect"
|
||||
)
|
||||
materialized_output = conv3d(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=conv3d,
|
||||
patch_fn=patched_module.torch_nn_conv3d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=conv3d,
|
||||
patch_fn=patched_module.torch_nn_conv3d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
|
||||
@clear_cache_before_run()
|
||||
@@ -267,21 +284,25 @@ def test_conv_transpose1d():
|
||||
|
||||
convtrans1d = torch.nn.ConvTranspose1d(in_channels=3, out_channels=4, kernel_size=2)
|
||||
materialized_output = convtrans1d(data)
|
||||
meta_data = data.to('meta')
|
||||
_assert_output_shape(data=meta_data,
|
||||
module=convtrans1d,
|
||||
patch_fn=patched_module.torch_nn_convtranspose1d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
meta_data = data.to("meta")
|
||||
_assert_output_shape(
|
||||
data=meta_data,
|
||||
module=convtrans1d,
|
||||
patch_fn=patched_module.torch_nn_convtranspose1d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
convtrans1d = torch.nn.ConvTranspose1d(in_channels=3, out_channels=4, kernel_size=2, padding=1)
|
||||
materialized_output = convtrans1d(data)
|
||||
meta_data = data.to('meta')
|
||||
_assert_output_shape(data=meta_data,
|
||||
module=convtrans1d,
|
||||
patch_fn=patched_module.torch_nn_convtranspose1d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
meta_data = data.to("meta")
|
||||
_assert_output_shape(
|
||||
data=meta_data,
|
||||
module=convtrans1d,
|
||||
patch_fn=patched_module.torch_nn_convtranspose1d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
|
||||
@clear_cache_before_run()
|
||||
@@ -291,21 +312,25 @@ def test_conv_transpose2d():
|
||||
|
||||
convtrans2d = torch.nn.ConvTranspose2d(in_channels=3, out_channels=4, kernel_size=2)
|
||||
materialized_output = convtrans2d(data)
|
||||
meta_data = data.to('meta')
|
||||
_assert_output_shape(data=meta_data,
|
||||
module=convtrans2d,
|
||||
patch_fn=patched_module.torch_nn_convtranspose2d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
meta_data = data.to("meta")
|
||||
_assert_output_shape(
|
||||
data=meta_data,
|
||||
module=convtrans2d,
|
||||
patch_fn=patched_module.torch_nn_convtranspose2d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
convtrans2d = torch.nn.ConvTranspose2d(in_channels=3, out_channels=4, kernel_size=2, padding=1)
|
||||
materialized_output = convtrans2d(data)
|
||||
meta_data = data.to('meta')
|
||||
_assert_output_shape(data=meta_data,
|
||||
module=convtrans2d,
|
||||
patch_fn=patched_module.torch_nn_convtranspose2d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
meta_data = data.to("meta")
|
||||
_assert_output_shape(
|
||||
data=meta_data,
|
||||
module=convtrans2d,
|
||||
patch_fn=patched_module.torch_nn_convtranspose2d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
|
||||
@clear_cache_before_run()
|
||||
@@ -315,46 +340,56 @@ def test_conv_transpose3d():
|
||||
|
||||
convtrans3d = torch.nn.ConvTranspose3d(in_channels=3, out_channels=4, kernel_size=2)
|
||||
materialized_output = convtrans3d(data)
|
||||
meta_data = data.to('meta')
|
||||
_assert_output_shape(data=meta_data,
|
||||
module=convtrans3d,
|
||||
patch_fn=patched_module.torch_nn_convtranspose3d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
meta_data = data.to("meta")
|
||||
_assert_output_shape(
|
||||
data=meta_data,
|
||||
module=convtrans3d,
|
||||
patch_fn=patched_module.torch_nn_convtranspose3d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
convtrans3d = torch.nn.ConvTranspose3d(in_channels=3, out_channels=4, kernel_size=2, padding=1)
|
||||
materialized_output = convtrans3d(data)
|
||||
meta_data = data.to('meta')
|
||||
_assert_output_shape(data=meta_data,
|
||||
module=convtrans3d,
|
||||
patch_fn=patched_module.torch_nn_convtranspose3d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
meta_data = data.to("meta")
|
||||
_assert_output_shape(
|
||||
data=meta_data,
|
||||
module=convtrans3d,
|
||||
patch_fn=patched_module.torch_nn_convtranspose3d,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
|
||||
@clear_cache_before_run()
|
||||
def test_pool1d():
|
||||
combinations = [[torch.nn.MaxPool1d, patched_module.torch_nn_maxpool1d],
|
||||
[torch.nn.AvgPool1d, patched_module.torch_nn_avgpool1d]]
|
||||
combinations = [
|
||||
[torch.nn.MaxPool1d, patched_module.torch_nn_maxpool1d],
|
||||
[torch.nn.AvgPool1d, patched_module.torch_nn_avgpool1d],
|
||||
]
|
||||
|
||||
for (layer_cls, patch_func) in combinations:
|
||||
for layer_cls, patch_func in combinations:
|
||||
pooler = layer_cls(kernel_size=3)
|
||||
|
||||
data = torch.rand(2, 3, 4)
|
||||
materialized_output = pooler(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=pooler,
|
||||
patch_fn=patch_func,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=pooler,
|
||||
patch_fn=patch_func,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
data = torch.rand(2, 4)
|
||||
materialized_output = pooler(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=pooler,
|
||||
patch_fn=patch_func,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=pooler,
|
||||
patch_fn=patch_func,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
data = torch.rand(2, 3, 4, 4)
|
||||
_assert_output_shape(data=data, module=pooler, patch_fn=patch_func, expect_exception=True, output_shape=None)
|
||||
@@ -362,29 +397,35 @@ def test_pool1d():
|
||||
|
||||
@clear_cache_before_run()
|
||||
def test_pool2d():
|
||||
combinations = [[torch.nn.MaxPool2d, patched_module.torch_nn_maxpool2d],
|
||||
[torch.nn.AvgPool2d, patched_module.torch_nn_avgpool2d]]
|
||||
combinations = [
|
||||
[torch.nn.MaxPool2d, patched_module.torch_nn_maxpool2d],
|
||||
[torch.nn.AvgPool2d, patched_module.torch_nn_avgpool2d],
|
||||
]
|
||||
|
||||
for (layer_cls, patch_func) in combinations:
|
||||
for layer_cls, patch_func in combinations:
|
||||
pooler = layer_cls(kernel_size=3)
|
||||
|
||||
# test max pool 3d
|
||||
data = torch.rand(2, 3, 4, 4)
|
||||
materialized_output = pooler(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=pooler,
|
||||
patch_fn=patch_func,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=pooler,
|
||||
patch_fn=patch_func,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
# test max pool 3d
|
||||
data = torch.rand(2, 4, 4)
|
||||
materialized_output = pooler(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=pooler,
|
||||
patch_fn=patch_func,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=pooler,
|
||||
patch_fn=patch_func,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
# test max pool 3d
|
||||
data = torch.rand(2, 3, 4, 4, 4)
|
||||
@@ -393,29 +434,35 @@ def test_pool2d():
|
||||
|
||||
@clear_cache_before_run()
|
||||
def test_pool3d():
|
||||
combinations = [[torch.nn.MaxPool3d, patched_module.torch_nn_maxpool3d],
|
||||
[torch.nn.AvgPool3d, patched_module.torch_nn_avgpool3d]]
|
||||
combinations = [
|
||||
[torch.nn.MaxPool3d, patched_module.torch_nn_maxpool3d],
|
||||
[torch.nn.AvgPool3d, patched_module.torch_nn_avgpool3d],
|
||||
]
|
||||
|
||||
for (layer_cls, patch_func) in combinations:
|
||||
for layer_cls, patch_func in combinations:
|
||||
pooler = layer_cls(kernel_size=3)
|
||||
|
||||
# test max pool 3d
|
||||
data = torch.rand(2, 3, 4, 4, 4)
|
||||
materialized_output = pooler(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=pooler,
|
||||
patch_fn=patch_func,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=pooler,
|
||||
patch_fn=patch_func,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
# test max pool 3d
|
||||
data = torch.rand(2, 4, 4, 4)
|
||||
materialized_output = pooler(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=pooler,
|
||||
patch_fn=patch_func,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape)
|
||||
_assert_output_shape(
|
||||
data=data,
|
||||
module=pooler,
|
||||
patch_fn=patch_func,
|
||||
expect_exception=False,
|
||||
output_shape=materialized_output.shape,
|
||||
)
|
||||
|
||||
# test max pool 3d
|
||||
data = torch.rand(2, 3, 4)
|
||||
@@ -430,19 +477,15 @@ def test_adaptive_pooling_1d():
|
||||
|
||||
data = torch.rand(3, 4)
|
||||
output = pooler(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=pooler,
|
||||
patch_fn=patch_func,
|
||||
expect_exception=False,
|
||||
output_shape=output.shape)
|
||||
_assert_output_shape(
|
||||
data=data, module=pooler, patch_fn=patch_func, expect_exception=False, output_shape=output.shape
|
||||
)
|
||||
|
||||
data = torch.rand(2, 3, 4)
|
||||
output = pooler(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=pooler,
|
||||
patch_fn=patch_func,
|
||||
expect_exception=False,
|
||||
output_shape=output.shape)
|
||||
_assert_output_shape(
|
||||
data=data, module=pooler, patch_fn=patch_func, expect_exception=False, output_shape=output.shape
|
||||
)
|
||||
|
||||
data = torch.rand(2, 3, 4, 5)
|
||||
_assert_output_shape(data=data, module=pooler, patch_fn=patch_func, expect_exception=True, output_shape=None)
|
||||
@@ -458,19 +501,15 @@ def test_adaptive_pooling_2d():
|
||||
|
||||
data = torch.rand(2, 3, 4)
|
||||
output = pooler(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=pooler,
|
||||
patch_fn=patch_func,
|
||||
expect_exception=False,
|
||||
output_shape=output.shape)
|
||||
_assert_output_shape(
|
||||
data=data, module=pooler, patch_fn=patch_func, expect_exception=False, output_shape=output.shape
|
||||
)
|
||||
|
||||
data = torch.rand(2, 3, 4, 5)
|
||||
output = pooler(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=pooler,
|
||||
patch_fn=patch_func,
|
||||
expect_exception=False,
|
||||
output_shape=output.shape)
|
||||
_assert_output_shape(
|
||||
data=data, module=pooler, patch_fn=patch_func, expect_exception=False, output_shape=output.shape
|
||||
)
|
||||
|
||||
|
||||
@clear_cache_before_run()
|
||||
@@ -483,16 +522,12 @@ def test_adaptive_pooling_3d():
|
||||
|
||||
data = torch.rand(2, 3, 4, 5)
|
||||
output = pooler(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=pooler,
|
||||
patch_fn=patch_func,
|
||||
expect_exception=False,
|
||||
output_shape=output.shape)
|
||||
_assert_output_shape(
|
||||
data=data, module=pooler, patch_fn=patch_func, expect_exception=False, output_shape=output.shape
|
||||
)
|
||||
|
||||
data = torch.rand(2, 3, 4, 5, 6)
|
||||
output = pooler(data)
|
||||
_assert_output_shape(data=data,
|
||||
module=pooler,
|
||||
patch_fn=patch_func,
|
||||
expect_exception=False,
|
||||
output_shape=output.shape)
|
||||
_assert_output_shape(
|
||||
data=data, module=pooler, patch_fn=patch_func, expect_exception=False, output_shape=output.shape
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user