diff --git a/test/test_schedule.py b/test/test_schedule.py index 71cec09068..67ab3fdbbe 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -970,17 +970,16 @@ class TestSchedule(unittest.TestCase): check_schedule(opt.schedule_step(), 9) def test_sgd_4convs_fuse(self): - with Context(FUSE_CONV_BW=1): - with Tensor.train(): - img = Tensor.empty(2,3,64,64) - c1 = nn.Conv2d(3,4,3,bias=False) - c2 = nn.Conv2d(4,8,3,bias=False) - c3 = nn.Conv2d(8,16,3,bias=False) - c4 = nn.Conv2d(16,32,3,bias=False) - opt = nn.optim.SGD(nn.state.get_parameters([c1, c2, c3, c4])) - opt.zero_grad() - c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward() - check_schedule(opt.schedule_step(), 19) + with Tensor.train(): + img = Tensor.empty(2,3,64,64) + c1 = nn.Conv2d(3,4,3,bias=False) + c2 = nn.Conv2d(4,8,3,bias=False) + c3 = nn.Conv2d(8,16,3,bias=False) + c4 = nn.Conv2d(16,32,3,bias=False) + opt = nn.optim.SGD(nn.state.get_parameters([c1, c2, c3, c4])) + opt.zero_grad() + c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward() + check_schedule(opt.schedule_step(), 22) @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") def test_prefer_half_buffer(self): @@ -1346,6 +1345,18 @@ class TestConvBW(unittest.TestCase): np.testing.assert_allclose(c1.weight.grad.numpy(), c1_torch.weight.grad.numpy(), atol=5e-4, rtol=1e-5) np.testing.assert_allclose(img.grad.numpy(), img_torch.grad.numpy(), atol=5e-4, rtol=1e-5) + def test_sgd_4convs_fuse(self): + with Tensor.train(): + img = Tensor.empty(2,3,64,64) + c1 = nn.Conv2d(3,4,3,bias=False) + c2 = nn.Conv2d(4,8,3,bias=False) + c3 = nn.Conv2d(8,16,3,bias=False) + c4 = nn.Conv2d(16,32,3,bias=False) + opt = nn.optim.SGD(nn.state.get_parameters([c1, c2, c3, c4])) + opt.zero_grad() + c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward() + self.check_schedule(opt.schedule_step(), 19) + class TestIndexing(unittest.TestCase): def check_schedule(self, xt:Union[Tensor,List[Tensor]], cnt:int): with Context(FUSE_ARANGE=getenv("FUSE_ARANGE", 1)):