def test_AvgPool2d(self):
from torch.nn import AvgPool2d as t_AvgPool2d
jt_model = AvgPool2d(3, 1, 1, ceil_mode=True)
torch_model = t_AvgPool2d(3, 1, 1, ceil_mode=True)
shape = (2, 16, 33, 33)
check(jt_model, torch_model, shape, False)
jt_model = AvgPool2d(3, 1, 1, ceil_mode=True, count_include_pad=False)
torch_model = t_AvgPool2d(3, 1, 1, ceil_mode=True, count_include_pad=False)
shape = (2, 16, 100, 100)
check(jt_model, torch_model, shape, False)
print('finish')
def test_cuda_old_pool(self):
from torch.nn import AvgPool2d
jt_model = OldPool(3, 1, 1, op="mean")
torch_model = AvgPool2d(3, 1, 1)
shape = [64, 64, 300, 300]
check(jt_model, torch_model, shape, False)
shape = [32, 128, 157, 300]
check(jt_model, torch_model, shape, False)
for i in range(10):
check(jt_model, torch_model, [1, 1, 300, 300], True)
def test_cpu_avg_pool2(self):
from torch.nn import AvgPool2d
jt_model = Pool(3, 1, 1, op="mean", ceil_mode=True)
torch_model = AvgPool2d(3, 1, 1, ceil_mode=True)
shape = (2, 16, 33, 33)
check(jt_model, torch_model, shape, False)