def test_save_and_load(self): state_dict = collections.OrderedDict([ ('a', torch.Tensor(2, 3)), ('b', 1), ('c', [1, 2, 3]), ('d', {'e': torch.Tensor(2), 'f': [4, 5, 6], 'g': {'h': 1}}), ]) f = io.BytesIO() torch.save(state_dict, f) f.seek(0) state_dict2 = torch.load(f) self.assertEqual(state_dict['b'], state_dict2['b']) self.assertEqual(state_dict['c'], state_dict2['c']) self.assertEqual(state_dict['d']['f'], state_dict2['d']['f']) self.assertEqual(state_dict['d']['g']['h'], state_dict2['d']['g']['h']) f = io.BytesIO() torch.save(torch.Tensor(2, 3), f) f = io.BytesIO() torch.save([1, 2, 3], f) f = '/tmp/test_dragon_vm_torch_save' try: torch.save(state_dict, f) state_dict2 = torch.load(f) self.assertEqual(state_dict['b'], state_dict2['b']) self.assertEqual(state_dict['c'], state_dict2['c']) self.assertEqual(state_dict['d']['f'], state_dict2['d']['f']) self.assertEqual(state_dict['d']['g']['h'], state_dict2['d']['g']['h']) except (OSError, PermissionError): pass
def test_variance_scaling(self): a = torch.Tensor(2, dtype=torch.float32) b = torch.Tensor(2, 3, dtype=torch.float32) c = torch.Tensor(2, 3, 3, dtype=torch.float32) entries = [('xavier_normal_', {}), ('xavier_uniform_', {}), ('kaiming_normal_', {}), ('kaiming_normal_', { 'nonlinearity': 'sigmoid' }), ('kaiming_normal_', { 'nonlinearity': 'tanh' }), ('kaiming_normal_', { 'nonlinearity': 'relu' }), ('kaiming_normal_', { 'nonlinearity': 'abc' }), ('kaiming_normal_', { 'a': None }), ('kaiming_normal_', { 'a': 'a' }), ('kaiming_normal_', { 'mode': 'fan_avg' }), ('kaiming_uniform_', {})] for init_name, kwargs in entries: try: getattr(torch.nn.init, init_name)(c, **kwargs) getattr(torch.nn.init, init_name)(b, **kwargs) getattr(torch.nn.init, init_name)(a, **kwargs) except ValueError: pass
def test_eye(self): x = torch.nn.init.eye_(torch.Tensor(2, 3, dtype=torch.float32)) self.assertEqual(x, np.eye(2, 3)) try: _ = torch.nn.init.eye_(torch.Tensor(2, 3, 3, dtype=torch.float32)) except ValueError: pass
def test_constant(self): x = torch.Tensor(2, dtype=torch.float32) self.assertEqual(torch.nn.init.constant_(x, 1), np.ones((2, ), dtype='float32')) self.assertEqual(torch.nn.init.ones_(x), np.ones((2, ), dtype='float32')) self.assertEqual(torch.nn.init.zeros_(x), np.zeros((2, ), dtype='float32'))
def test_dirac(self): for ndim in range(2, 6): for groups in range(1, 4): try: _ = torch.nn.init.dirac_(torch.Tensor(*([2] * ndim), dtype=torch.float32), groups=groups) except ValueError: pass
class TestJit(unittest.TestCase): """Test the jit component.""" @torch.jit.trace(example_inputs=[ torch.Tensor(1, dtype=torch.int64), torch.Tensor(1, dtype=torch.int64), ]) def func1(self, a, b, **kwargs): _ = kwargs return a + b def test_trace(self): @torch.jit.trace(example_inputs=[None, None]) def func2(a, b): return a + b @torch.jit.trace def func3(a, b): return a + b @torch.jit.trace(example_inputs=[None]) def func4(a, b): return a + b class TestModule(torch.nn.Module): def forward(self, a, b): return a + b func5 = torch.jit.trace(lambda a, b: a + b) m = torch.jit.trace(TestModule()) a, b, c = torch.tensor([1, 2]), torch.tensor([3, 4]), torch.tensor(1) self.assertEqual(self.func1(a, b).numpy().tolist(), [4, 6]) self.assertEqual(func2(a, b).numpy().tolist(), [4, 6]) self.assertEqual(func3(a, b).numpy().tolist(), [4, 6]) self.assertEqual(func5(a, b).numpy().tolist(), [4, 6]) self.assertEqual(m(a, b).numpy().tolist(), [4, 6]) self.assertEqual(self.func1(a, b, c=c).numpy().tolist(), [4, 6]) try: func4(a, b) except ValueError: pass
def test_properties(self): a = torch.tensor([0.]).cpu() b = torch.Tensor([0., 1.], dtype=torch.float64).zero_() a.requires_grad = True c = a + b c.retain_grad() c.backward() self.assertEqual(a.is_leaf, True) self.assertEqual(a.is_floating_point(), True) self.assertEqual(a.is_contiguous(), True) self.assertEqual(a.contiguous().is_contiguous(), True) self.assertEqual(a.volatile, False) self.assertEqual(a.numel(), 1) self.assertEqual(a.grad_fn, None) self.assertEqual(float(a.grad), 2.) self.assertEqual(b.grad, None) self.assertEqual(int(a.detach()), 0) self.assertEqual(torch.Tensor([0]).dim(), 1) self.assertEqual(float(torch.Tensor(1).one_()), 1.) self.assertEqual(torch.tensor(2.333).item(), 2.333) self.assertEqual(torch.tensor([2, 3]).tolist(), [2, 3]) self.assertEqual(torch.empty(2, 3).ndimension(), 2) self.assertEqual(torch.empty(3).new_empty(2, 3).ndimension(), 2) self.assertEqual(repr(torch.tensor(1)), '1') self.assertEqual(repr(torch.tensor(1).new_tensor(1)), '1') self.assertNotEqual(a.__hash__(), b.__hash__()) self.assertNotEqual(a.__repr__(), b.__repr__()) self.assertEqual(torch.BoolTensor(1).dtype, 'bool') self.assertEqual(torch.ByteTensor(1).dtype, 'uint8') self.assertEqual(torch.CharTensor(1).dtype, 'int8') self.assertEqual(torch.DoubleTensor(1).dtype, 'float64') self.assertEqual(torch.FloatTensor(1).dtype, 'float32') self.assertEqual(torch.HalfTensor(1).dtype, 'float16') self.assertEqual(torch.IntTensor(1).dtype, 'int32') self.assertEqual(torch.LongTensor(1).dtype, 'int64') self.assertEqual(torch.autograd.Variable(torch.Tensor(1)).requires_grad, False) try: _ = torch.Tensor(5.) except ValueError: pass try: _ = torch.Tensor(2, 3.) except ValueError: pass try: torch.Tensor(2).retain_grad() except RuntimeError: pass
def new_tensor(shape, dtype, device): """Return a new tensor abstraction.""" return torch.Tensor(*shape, dtype=dtype, device=device)
def test_random(self): _ = torch.nn.init.normal_(torch.Tensor(2, dtype=torch.float32)) _ = torch.nn.init.uniform_(torch.Tensor(2, dtype=torch.float32))
def test_constant(self): x = torch.nn.init.constant_(torch.Tensor(2, dtype=torch.float32), 1) self.assertEqual(x, np.ones((2, ), dtype='float32'))