mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-10 23:48:01 -05:00
int8/uint8 support (#837)
* feat: int8 support * feat: uint8 support * feat: int8 tests * fix: fix uint8 on clang * feat: test casting between int8/uint8/float16/float32 * clean: way cleaner dtype tests * feat: preprocess_imagenet using the correct dtype * feat: add test for overflow between uint8 and int8
This commit is contained in:
wozeparrot
@@ -8,83 +8,110 @@ from tinygrad.tensor import Tensor, dtypes
|
||||
# for LLVM, it segfaults because it can't link to the casting function
|
||||
@unittest.skipIf(getenv("CI", "") != "" and Device.DEFAULT in ["LLVM"], "float16 broken in some CI backends")
|
||||
class TestDtype(unittest.TestCase):
|
||||
def test_half_to_np(self):
|
||||
a = Tensor([1,2,3,4], dtype=dtypes.float16)
|
||||
def _test_to_np(self, a, np_dtype, target):
|
||||
print(a)
|
||||
na = a.numpy()
|
||||
print(na, na.dtype, a.lazydata.realized)
|
||||
assert na.dtype == np.float16
|
||||
np.testing.assert_allclose(na, [1,2,3,4])
|
||||
assert na.dtype == np_dtype
|
||||
np.testing.assert_allclose(na, target)
|
||||
|
||||
def test_half_add(self):
|
||||
a = Tensor([1,2,3,4], dtype=dtypes.float16)
|
||||
b = Tensor([1,2,3,4], dtype=dtypes.float16)
|
||||
def test_half_to_np(self): self._test_to_np(Tensor([1,2,3,4], dtype=dtypes.float16), np.float16, [1,2,3,4])
|
||||
def test_int8_to_np(self): self._test_to_np(Tensor([1,2,3,4], dtype=dtypes.int8), np.int8, [1,2,3,4])
|
||||
def test_uint8_to_np(self): self._test_to_np(Tensor([1,2,3,4], dtype=dtypes.uint8), np.uint8, [1,2,3,4])
|
||||
|
||||
def _test_cast(self, a, target_dtype, target):
|
||||
print(a)
|
||||
b = a.cast(target_dtype)
|
||||
print(b.numpy())
|
||||
assert b.dtype == target_dtype
|
||||
np.testing.assert_allclose(b.numpy(), target)
|
||||
|
||||
def test_float_to_half(self): self._test_cast(Tensor([1,2,3,4], dtype=dtypes.float32), dtypes.float16, [1,2,3,4])
|
||||
def test_float_to_int8(self): self._test_cast(Tensor([1,2,3,4], dtype=dtypes.float32), dtypes.int8, [1,2,3,4])
|
||||
def test_float_to_uint8(self): self._test_cast(Tensor([1,2,3,4], dtype=dtypes.float32), dtypes.uint8, [1,2,3,4])
|
||||
|
||||
def test_half_to_float(self): self._test_cast(Tensor([1,2,3,4], dtype=dtypes.float16), dtypes.float32, [1,2,3,4])
|
||||
def test_half_to_int8(self): self._test_cast(Tensor([1,2,3,4], dtype=dtypes.float16), dtypes.int8, [1,2,3,4])
|
||||
def test_half_to_uint8(self): self._test_cast(Tensor([1,2,3,4], dtype=dtypes.float16), dtypes.uint8, [1,2,3,4])
|
||||
|
||||
def test_int8_to_float(self): self._test_cast(Tensor([1,2,3,4], dtype=dtypes.int8), dtypes.float32, [1,2,3,4])
|
||||
def test_int8_to_half(self): self._test_cast(Tensor([1,2,3,4], dtype=dtypes.int8), dtypes.float16, [1,2,3,4])
|
||||
def test_int8_to_uint8(self): self._test_cast(Tensor([1,2,3,4], dtype=dtypes.int8), dtypes.uint8, [1,2,3,4])
|
||||
|
||||
def test_uint8_to_float(self): self._test_cast(Tensor([1,2,3,4], dtype=dtypes.uint8), dtypes.float32, [1,2,3,4])
|
||||
def test_uint8_to_half(self): self._test_cast(Tensor([1,2,3,4], dtype=dtypes.uint8), dtypes.float16, [1,2,3,4])
|
||||
def test_uint8_to_int8(self): self._test_cast(Tensor([1,2,3,4], dtype=dtypes.uint8), dtypes.int8, [1,2,3,4])
|
||||
|
||||
def _test_add(self, a, b, target_dtype, target):
|
||||
c = a+b
|
||||
print(c.numpy())
|
||||
assert c.dtype == dtypes.float16
|
||||
np.testing.assert_allclose(c.numpy(), [2,4,6,8])
|
||||
assert c.dtype == target_dtype
|
||||
np.testing.assert_allclose(c.numpy(), target)
|
||||
|
||||
def test_half_mul(self):
|
||||
a = Tensor([1,2,3,4], dtype=dtypes.float16)
|
||||
b = Tensor([1,2,3,4], dtype=dtypes.float16)
|
||||
def test_half_add(self): self._test_add(Tensor([1,2,3,4], dtype=dtypes.float16), Tensor([1,2,3,4], dtype=dtypes.float16), dtypes.float16, [2,4,6,8])
|
||||
def test_int8_add(self): self._test_add(Tensor([1,2,3,4], dtype=dtypes.int8), Tensor([1,2,3,4], dtype=dtypes.int8), dtypes.int8, [2,4,6,8])
|
||||
|
||||
def _test_mul(self, a, b, target_dtype, target):
|
||||
c = a*b
|
||||
print(c.numpy())
|
||||
assert c.dtype == dtypes.float16
|
||||
np.testing.assert_allclose(c.numpy(), [1,4,9,16])
|
||||
assert c.dtype == target_dtype
|
||||
np.testing.assert_allclose(c.numpy(), target)
|
||||
|
||||
def test_half_matmul(self):
|
||||
a = Tensor([[1,2],[3,4]], dtype=dtypes.float16)
|
||||
b = Tensor.eye(2, dtype=dtypes.float16)
|
||||
def test_half_mul(self): self._test_mul(Tensor([1,2,3,4], dtype=dtypes.float16), Tensor([1,2,3,4], dtype=dtypes.float16), dtypes.float16, [1,4,9,16])
|
||||
def test_int8_mul(self): self._test_mul(Tensor([1,2,3,4], dtype=dtypes.int8), Tensor([1,2,3,4], dtype=dtypes.int8), dtypes.int8, [1,4,9,16])
|
||||
|
||||
def _test_matmul(self, a, b, target_dtype, target):
|
||||
c = a@b
|
||||
print(c.numpy())
|
||||
assert c.dtype == dtypes.float16
|
||||
np.testing.assert_allclose(c.numpy(), [[1,2],[3,4]])
|
||||
assert c.dtype == target_dtype
|
||||
np.testing.assert_allclose(c.numpy(), target)
|
||||
|
||||
def test_upcast_float(self):
|
||||
a = Tensor([1,2,3,4], dtype=dtypes.float16)
|
||||
print(a)
|
||||
fa = a.float()
|
||||
assert a.device == fa.device
|
||||
assert a.requires_grad == fa.requires_grad
|
||||
na = fa.numpy()
|
||||
print(na, na.dtype)
|
||||
assert na.dtype == np.float32
|
||||
np.testing.assert_allclose(na, [1,2,3,4])
|
||||
def test_half_matmul(self): self._test_matmul(Tensor([[1,2],[3,4]], dtype=dtypes.float16), Tensor.eye(2, dtype=dtypes.float16), dtypes.float16, [[1,2],[3,4]])
|
||||
def test_int8_matmul(self): self._test_matmul(Tensor([[1,2],[3,4]], dtype=dtypes.int8), Tensor.eye(2, dtype=dtypes.int8), dtypes.int8, [[1,2],[3,4]])
|
||||
|
||||
def test_downcast_float(self):
|
||||
a = Tensor([1,2,3,4], dtype=dtypes.float32, requires_grad=False).half()
|
||||
print(a)
|
||||
ha = a.half()
|
||||
assert a.device == ha.device
|
||||
assert a.requires_grad == ha.requires_grad
|
||||
na = ha.numpy()
|
||||
print(na, na.dtype)
|
||||
assert na.dtype == np.float16
|
||||
np.testing.assert_allclose(na, [1,2,3,4])
|
||||
|
||||
def test_half_add_upcast(self):
|
||||
a = Tensor([1,2,3,4], dtype=dtypes.float16)
|
||||
b = Tensor([1,2,3,4], dtype=dtypes.float32)
|
||||
def _test_add_upcast(self, a, b, target_dtype, target):
|
||||
c = a+b
|
||||
print(c.numpy())
|
||||
assert c.dtype == dtypes.float32
|
||||
np.testing.assert_allclose(c.numpy(), [2,4,6,8])
|
||||
assert c.dtype == target_dtype
|
||||
np.testing.assert_allclose(c.numpy(), target)
|
||||
|
||||
def test_half_mul_upcast(self):
|
||||
a = Tensor([1,2,3,4], dtype=dtypes.float16)
|
||||
b = Tensor([1,2,3,4], dtype=dtypes.float32)
|
||||
def test_half_add_upcast_float(self): self._test_add_upcast(Tensor([1,2,3,4], dtype=dtypes.float16), Tensor([1,2,3,4], dtype=dtypes.float32), dtypes.float32, [2,4,6,8])
|
||||
def test_int8_add_upcast_float(self): self._test_add_upcast(Tensor([1,2,3,4], dtype=dtypes.int8), Tensor([1,2,3,4], dtype=dtypes.float32), dtypes.float32, [2,4,6,8])
|
||||
def test_int8_add_upcast_half(self): self._test_add_upcast(Tensor([1,2,3,4], dtype=dtypes.int8), Tensor([1,2,3,4], dtype=dtypes.float16), dtypes.float16, [2,4,6,8])
|
||||
|
||||
def _test_mul_upcast(self, a, b, target_dtype, target):
|
||||
c = a*b
|
||||
print(c.numpy())
|
||||
assert c.dtype == dtypes.float32
|
||||
np.testing.assert_allclose(c.numpy(), [1,4,9,16])
|
||||
assert c.dtype == target_dtype
|
||||
np.testing.assert_allclose(c.numpy(), target)
|
||||
|
||||
def test_half_matmul_upcast(self):
|
||||
a = Tensor([[1,2],[3,4]], dtype=dtypes.float16)
|
||||
b = Tensor.eye(2, dtype=dtypes.float32)
|
||||
def test_half_mul_upcast_float(self): self._test_mul_upcast(Tensor([1,2,3,4], dtype=dtypes.float16), Tensor([1,2,3,4], dtype=dtypes.float32), dtypes.float32, [1,4,9,16])
|
||||
def test_int8_mul_upcast_float(self): self._test_mul_upcast(Tensor([1,2,3,4], dtype=dtypes.int8), Tensor([1,2,3,4], dtype=dtypes.float32), dtypes.float32, [1,4,9,16])
|
||||
def test_int8_mul_upcast_half(self): self._test_mul_upcast(Tensor([1,2,3,4], dtype=dtypes.int8), Tensor([1,2,3,4], dtype=dtypes.float16), dtypes.float16, [1,4,9,16])
|
||||
|
||||
def _test_matmul_upcast(self, a, b, target_dtype, target):
|
||||
c = a@b
|
||||
print(c.numpy())
|
||||
assert c.dtype == dtypes.float32
|
||||
np.testing.assert_allclose(c.numpy(), [[1,2],[3,4]])
|
||||
assert c.dtype == target_dtype
|
||||
np.testing.assert_allclose(c.numpy(), target)
|
||||
|
||||
def test_half_matmul_upcast_float(self): self._test_matmul_upcast(Tensor([[1,2],[3,4]], dtype=dtypes.float16), Tensor.eye(2, dtype=dtypes.float32), dtypes.float32, [[1,2],[3,4]])
|
||||
def test_int8_matmul_upcast_float(self): self._test_matmul_upcast(Tensor([[1,2],[3,4]], dtype=dtypes.int8), Tensor.eye(2, dtype=dtypes.float32), dtypes.float32, [[1,2],[3,4]])
|
||||
def test_int8_matmul_upcast_half(self): self._test_matmul_upcast(Tensor([[1,2],[3,4]], dtype=dtypes.int8), Tensor.eye(2, dtype=dtypes.float16), dtypes.float16, [[1,2],[3,4]])
|
||||
|
||||
def test_int8_to_uint8_negative(self):
|
||||
a = Tensor([-1, -2, -3, -4], dtype=dtypes.int8)
|
||||
print(a)
|
||||
b = a.cast(dtypes.uint8)
|
||||
print(b.numpy())
|
||||
np.testing.assert_allclose(b.numpy(), [255, 254, 253, 252])
|
||||
|
||||
def test_uint8_to_int8_overflow(self):
|
||||
a = Tensor([255, 254, 253, 252], dtype=dtypes.uint8)
|
||||
print(a)
|
||||
b = a.cast(dtypes.int8)
|
||||
print(b.numpy())
|
||||
np.testing.assert_allclose(b.numpy(), [-1, -2, -3, -4])
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user