mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
44436947bc
This reverts commit 8ef6356f267c75276ea23b51163274cd5fffc0ce. Reverted https://github.com/pytorch/pytorch/pull/79617 on behalf of https://github.com/zengk95 due to this is breaking periodic jobs (and maybe pull) on trunk
125 lines
3.8 KiB
Python
125 lines
3.8 KiB
Python
# Owner(s): ["oncall: jit"]
|
|
|
|
import torch
|
|
import torch._lazy
|
|
import torch._lazy.config
|
|
import torch._lazy.ir_cache
|
|
import torch._lazy.ts_backend
|
|
import torch._lazy.metrics as metrics
|
|
from torch.testing._internal.common_utils import IS_WINDOWS, run_tests, TestCase
|
|
import os
|
|
import unittest
|
|
|
|
torch._lazy.ts_backend.init()
|
|
torch._lazy.config.set_reuse_ir(True)
|
|
|
|
def get_test_device():
|
|
return 'cuda' if 'LTC_TS_CUDA' in os.environ else 'cpu'
|
|
|
|
@unittest.skipIf(IS_WINDOWS, "To be fixed")
|
|
class TestLazyReuseIr(TestCase):
|
|
def testAdd(self):
|
|
device = get_test_device()
|
|
x = torch.randn(2, 3, 4, device=device)
|
|
y = torch.randn(2, 3, 4, device=device)
|
|
z = torch.zeros(2, 3, 4, device=device)
|
|
|
|
device = 'lazy'
|
|
x_lazy = x.detach().clone().to(device=device)
|
|
y_lazy = y.detach().clone().to(device=device)
|
|
z_lazy = z.detach().clone().to(device=device)
|
|
|
|
for i in range(10):
|
|
z += (x + y)
|
|
|
|
for i in range(10):
|
|
z_lazy += (x_lazy + y_lazy)
|
|
torch._lazy.mark_step()
|
|
|
|
torch.testing.assert_close(z.cpu(), z_lazy.cpu())
|
|
assert metrics.counter_value("IrNodeReused_torch::lazy::AddTensor") >= 14
|
|
metrics.reset()
|
|
torch._lazy.ir_cache.reset()
|
|
|
|
def testAddSub(self):
|
|
device = get_test_device()
|
|
x = torch.randn(2, 3, 4, device=device)
|
|
y = torch.randn(2, 3, 4, device=device)
|
|
z = torch.zeros(2, 3, 4, device=device)
|
|
|
|
device = 'lazy'
|
|
x_lazy = x.detach().clone().to(device=device)
|
|
y_lazy = y.detach().clone().to(device=device)
|
|
z_lazy = z.detach().clone().to(device=device)
|
|
|
|
for i in range(10):
|
|
if i < 5:
|
|
z += (x + y)
|
|
else:
|
|
z += (x - y)
|
|
|
|
for i in range(10):
|
|
if i < 5:
|
|
z_lazy += (x_lazy + y_lazy)
|
|
else:
|
|
z_lazy += (x_lazy - y_lazy)
|
|
torch._lazy.mark_step()
|
|
|
|
torch.testing.assert_close(z.cpu(), z_lazy.cpu())
|
|
assert metrics.counter_value("IrNodeReused_torch::lazy::AddTensor") >= 8
|
|
metrics.reset()
|
|
torch._lazy.ir_cache.reset()
|
|
|
|
def testAddSubFallback(self):
|
|
torch._lazy.config.set_force_fallback("aten::sub")
|
|
device = get_test_device()
|
|
x = torch.randn(2, 3, 4, device=device)
|
|
y = torch.randn(2, 3, 4, device=device)
|
|
z = torch.zeros(2, 3, 4, device=device)
|
|
|
|
device = 'lazy'
|
|
x_lazy = x.detach().clone().to(device=device)
|
|
y_lazy = y.detach().clone().to(device=device)
|
|
z_lazy = z.detach().clone().to(device=device)
|
|
|
|
for i in range(10):
|
|
if i < 5:
|
|
z += (x + y)
|
|
else:
|
|
z += (x - y)
|
|
|
|
for i in range(10):
|
|
if i < 5:
|
|
z_lazy += (x_lazy + y_lazy)
|
|
else:
|
|
z_lazy += (x_lazy - y_lazy)
|
|
torch._lazy.mark_step()
|
|
|
|
torch.testing.assert_close(z.cpu(), z_lazy.cpu())
|
|
assert metrics.counter_value("IrNodeReused_torch::lazy::AddTensor") >= 8
|
|
metrics.reset()
|
|
torch._lazy.ir_cache.reset()
|
|
torch._lazy.config.set_force_fallback("")
|
|
|
|
def testBatchNorm(self):
|
|
device = get_test_device()
|
|
x = torch.randn(16, 3, 224, 224, device=device)
|
|
bn = torch.nn.BatchNorm2d(3).to(device=device)
|
|
for i in range(10):
|
|
z = bn(x)
|
|
|
|
device = "lazy"
|
|
x_lazy = x.detach().clone().to(device=device)
|
|
bn = bn.to(device=device)
|
|
for i in range(10):
|
|
z_lazy = bn(x_lazy)
|
|
torch._lazy.mark_step()
|
|
|
|
torch.testing.assert_close(z.cpu(), z_lazy.cpu())
|
|
assert metrics.counter_value("IrNodeReused_torch::lazy::TSNativeBatchNormForward") >= 7
|
|
metrics.reset()
|
|
torch._lazy.ir_cache.reset()
|
|
|
|
if __name__ == '__main__':
|
|
run_tests()
|