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'))
