wehub-resource-sync
176 lines
6.3 KiB
Python
176 lines
6.3 KiB
Python
import unittest
|
|
|
|
import numpy as np
|
|
import torch.cuda
|
|
|
|
import ray
|
|
from ray.rllib.utils.test_utils import check
|
|
from ray.rllib.utils.torch_utils import (
|
|
clip_gradients,
|
|
convert_to_torch_tensor,
|
|
copy_torch_tensors,
|
|
two_hot,
|
|
)
|
|
|
|
|
|
class TestTorchUtils(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls) -> None:
|
|
ray.init()
|
|
|
|
@classmethod
|
|
def tearDownClass(cls) -> None:
|
|
ray.shutdown()
|
|
|
|
def test_convert_to_torch_tensor(self):
|
|
# Tests whether convert_to_torch_tensor works as expected
|
|
|
|
# Test None
|
|
self.assertTrue(convert_to_torch_tensor(None) is None)
|
|
|
|
# Test single array
|
|
array = np.array([1, 2, 3])
|
|
tensor = torch.from_numpy(array)
|
|
self.assertTrue(all(convert_to_torch_tensor(array) == tensor))
|
|
|
|
# Test torch tensor
|
|
self.assertTrue(convert_to_torch_tensor(tensor) is tensor)
|
|
|
|
# Test conversion to 32-bit float
|
|
tensor_2 = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64)
|
|
self.assertTrue(convert_to_torch_tensor(tensor_2).dtype is torch.float32)
|
|
|
|
# Test nested structure with objects tested above
|
|
converted = convert_to_torch_tensor(
|
|
{"a": (array, tensor), "b": tensor_2, "c": None}
|
|
)
|
|
self.assertTrue(all(convert_to_torch_tensor(converted["a"][0]) == tensor))
|
|
self.assertTrue(converted["a"][1] is tensor)
|
|
self.assertTrue(converted["b"].dtype is torch.float32)
|
|
self.assertTrue(converted["c"] is None)
|
|
|
|
def test_copy_torch_tensors(self):
|
|
array = np.array([1, 2, 3], dtype=np.float32)
|
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
|
tensor = torch.from_numpy(array).to(device)
|
|
tensor_2 = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64).to(device)
|
|
|
|
# Test single tensor
|
|
copied_tensor = copy_torch_tensors(tensor, device)
|
|
self.assertTrue(copied_tensor.device == device)
|
|
self.assertNotEqual(id(copied_tensor), id(tensor))
|
|
self.assertTrue(all(copied_tensor == tensor))
|
|
|
|
# check that dtypes aren't modified
|
|
copied_tensor_2 = copy_torch_tensors(tensor_2, device)
|
|
self.assertTrue(copied_tensor_2.dtype == tensor_2.dtype)
|
|
self.assertFalse(copied_tensor_2.dtype == torch.float32)
|
|
|
|
# Test nested structure can be converted
|
|
nested_structure = {"a": tensor, "b": tensor_2, "c": 1}
|
|
copied_nested_structure = copy_torch_tensors(nested_structure, device)
|
|
self.assertTrue(copied_nested_structure["a"].device == device)
|
|
self.assertTrue(copied_nested_structure["b"].device == device)
|
|
self.assertTrue(copied_nested_structure["c"] == 1)
|
|
self.assertNotEqual(id(copied_nested_structure["a"]), id(tensor))
|
|
self.assertNotEqual(id(copied_nested_structure["b"]), id(tensor_2))
|
|
self.assertTrue(all(copied_nested_structure["a"] == tensor))
|
|
self.assertTrue(all(copied_nested_structure["b"] == tensor_2))
|
|
|
|
# if gpu is available test moving tensor from cpu to gpu and vice versa
|
|
if torch.cuda.is_available():
|
|
tensor = torch.from_numpy(array).to("cpu")
|
|
copied_tensor = copy_torch_tensors(tensor, "cuda:0")
|
|
self.assertFalse(copied_tensor.device == torch.device("cpu"))
|
|
self.assertTrue(copied_tensor.device == torch.device("cuda:0"))
|
|
self.assertNotEqual(id(copied_tensor), id(tensor))
|
|
self.assertTrue(
|
|
all(copied_tensor.detach().cpu().numpy() == tensor.detach().numpy())
|
|
)
|
|
|
|
tensor = torch.from_numpy(array).to("cuda:0")
|
|
copied_tensor = copy_torch_tensors(tensor, "cpu")
|
|
self.assertFalse(copied_tensor.device == torch.device("cuda:0"))
|
|
self.assertTrue(copied_tensor.device == torch.device("cpu"))
|
|
self.assertNotEqual(id(copied_tensor), id(tensor))
|
|
self.assertTrue(
|
|
all(copied_tensor.detach().numpy() == tensor.detach().cpu().numpy())
|
|
)
|
|
|
|
def test_large_gradients_clipping(self):
|
|
|
|
large_gradients = {
|
|
f"gradient_{i}": torch.full((256, 256), 1e22) for i in range(20)
|
|
}
|
|
|
|
total_norm = clip_gradients(
|
|
large_gradients, grad_clip=40, grad_clip_by="global_norm"
|
|
)
|
|
|
|
self.assertFalse(total_norm.isinf())
|
|
print(f"total norm for large gradients: {total_norm}")
|
|
|
|
small_gradients = {
|
|
f"gradient_{i}": torch.full((256, 256), 1e-22) for i in range(20)
|
|
}
|
|
|
|
total_norm = clip_gradients(
|
|
small_gradients, grad_clip=40, grad_clip_by="global_norm"
|
|
)
|
|
self.assertFalse(total_norm.isneginf())
|
|
print(f"total norm for small gradients: {total_norm}")
|
|
|
|
def test_two_hot(self):
|
|
# Test value that's exactly on one of the bucket boundaries. This used to return
|
|
# a two-hot vector with a NaN in it, as k == kp1 at that boundary.
|
|
check(
|
|
two_hot(torch.tensor([0.0]), 10, -5.0, 5.0),
|
|
np.array([[0, 0, 0, 0, 0.5, 0.5, 0, 0, 0, 0]]),
|
|
)
|
|
|
|
# Test violating the boundaries (upper and lower).
|
|
upper_bound = np.zeros((255,))
|
|
upper_bound[-1] = 1.0
|
|
lower_bound = np.zeros((255,))
|
|
lower_bound[0] = 1.0
|
|
check(
|
|
two_hot(torch.tensor([20.1, 50.0, 150.0, -20.00001])),
|
|
np.array([upper_bound, upper_bound, upper_bound, lower_bound]),
|
|
)
|
|
|
|
# Test other cases.
|
|
check(
|
|
two_hot(torch.tensor([2.5]), 11, -5.0, 5.0),
|
|
np.array([[0, 0, 0, 0, 0, 0, 0, 0.5, 0.5, 0, 0]]),
|
|
)
|
|
check(
|
|
two_hot(torch.tensor([2.5, 0.1]), 10, -5.0, 5.0),
|
|
np.array(
|
|
[
|
|
[0, 0, 0, 0, 0, 0, 0.25, 0.75, 0, 0],
|
|
[0, 0, 0, 0, 0.41, 0.59, 0, 0, 0, 0],
|
|
]
|
|
),
|
|
)
|
|
check(
|
|
two_hot(torch.tensor([0.1]), 4, -1.0, 1.0),
|
|
np.array([[0, 0.35, 0.65, 0]]),
|
|
)
|
|
check(
|
|
two_hot(torch.tensor([-0.5, -1.2]), 9, -6.0, 3.0),
|
|
np.array(
|
|
[
|
|
[0, 0, 0, 0, 0.11111, 0.88889, 0, 0, 0],
|
|
[0, 0, 0, 0, 0.73333, 0.26667, 0, 0, 0],
|
|
]
|
|
),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|