def test_link_as_torch_model():
# initialized parameter
a_arr = numpy.ones((3, 2), 'float32')
a_chainer_param = chainer.Parameter(a_arr)
# 0-size parameter
b_arr = numpy.ones((2, 0, 1), 'float32')
b_chainer_param = chainer.Parameter(b_arr)
link = chainer.Link()
with link.init_scope():
link.a = a_chainer_param
link.b = b_chainer_param
# Conversion
torched = cpm.LinkAsTorchModel(link)
params = list(torched.parameters())
assert len(params) == 2
assert isinstance(params[0], torch.nn.Parameter)
assert isinstance(params[1], torch.nn.Parameter)
assert params[0].shape == (3, 2)
assert params[1].shape == (2, 0, 1)
assert (params[0].data.numpy() == numpy.ones((3, 2))).all()
# Test memory sharing
params[0].data[...] = torch.tensor(numpy.arange(6).reshape((3, 2)))
assert (a_chainer_param.array == numpy.arange(6).reshape((3, 2))).all()
def test_link_as_torch_model_uninitialized():
# Uninitialized parameters are not supported
a_chainer_param = chainer.Parameter()
link = chainer.Link()
with link.init_scope():
link.a = a_chainer_param
with pytest.raises(TypeError):
torched = cpm.LinkAsTorchModel(link)
torched.parameters()
def set_model(self):
chainer_model = MLP()
chainer_model.to_device(self.device)
self.device.use()
dummy_input = self.train_dataset[0][0]
dummy_input = chainer.Variable(tensor.asarray(dummy_input))
dummy_input.to_device(self.device)
chainer_model(dummy_input)
# dummy_input = iter(self.train_loader).next()
# dummy_input = chainer_prepare_batch(dummy_input, self.device)
# chainer_model(dummy_input[0][0])
self.model = cpm.LinkAsTorchModel(chainer_model)
def test_named_params():
a_arr = numpy.ones((3, 2), 'float32')
a_chainer_param = chainer.Parameter(a_arr)
# 0-size parameter
b_arr = numpy.ones((2, 0, 1), 'float32')
b_chainer_param = chainer.Parameter(b_arr)
link = chainer.Link()
with link.init_scope():
link.a = a_chainer_param
link.b = b_chainer_param
torched = cpm.LinkAsTorchModel(link)
n_params = dict(torched.named_parameters())
assert 'a' in n_params
numpy.testing.assert_array_equal(a_arr, n_params['a'].detach())
assert 'b' in n_params
numpy.testing.assert_array_equal(b_arr, n_params['b'].detach())
def test_state_dict():
a_arr = numpy.ones((3, 2), 'float32')
a_chainer_param = chainer.Parameter(a_arr)
# 0-size parameter
b_arr = numpy.ones((2, 0, 1), 'float32')
b_chainer_param = chainer.Parameter(b_arr)
link = chainer.Link()
with link.init_scope():
link.a = a_chainer_param
link.b = b_chainer_param
torched = cpm.LinkAsTorchModel(link)
state_dict = torched.state_dict()
assert '/a' in state_dict
numpy.testing.assert_array_equal(a_arr, state_dict['/a'].detach())
assert '/b' in state_dict
numpy.testing.assert_array_equal(b_arr, state_dict['/b'].detach())
def test_link_to_device():
a_arr = numpy.ones((3, 2), 'float32')
a_chainer_param = chainer.Parameter(a_arr)
# 0-size parameter
b_arr = numpy.ones((2, 0, 1), 'float32')
b_chainer_param = chainer.Parameter(b_arr)
link = chainer.Link()
with link.init_scope():
link.a = a_chainer_param
link.b = b_chainer_param
torched = cpm.LinkAsTorchModel(link)
ret = torched.to('cuda')
assert torched is ret
for name, param in torched.named_parameters():
assert param.device.type == 'cuda'
def test_link_as_torch_model_nested():
dtype = numpy.float32
# l2: MyLink2 := p2 * l1(x)
# - p2
# - l1: MyLink1 := p1 * x
# - p1
class MyLink1(chainer.Link):
def __init__(self):
super().__init__()
with self.init_scope():
self.p1 = chainer.Parameter(numpy.array([2], dtype))
def forward(self, x1):
return self.p1 * x1
class MyLink2(chainer.Chain):
def __init__(self):
super().__init__()
with self.init_scope():
self.p2 = chainer.Parameter(numpy.array([3], dtype))
self.l1 = MyLink1()
def forward(self, x2):
return self.p2 * self.l1(x2)
# Dummy optimizer that always writes a constant value.
class MyOptim(torch.optim.Optimizer):
def __init__(self, params):
super().__init__(params, {})
self.constant = None
def set_constant(self, constant):
self.constant = constant
def step(self, closure=None):
for group in self.param_groups:
for param in group['params']:
param.data[...] = self.constant
link = MyLink2()
module = cpm.LinkAsTorchModel(link)
assert isinstance(module.p2, torch.nn.Parameter)
assert isinstance(module.l1, torch.nn.Module)
assert isinstance(module.l1.p1, torch.nn.Parameter)
assert len(list(module.parameters(recurse=False))) == 1
assert len(list(module.l1.parameters(recurse=False))) == 1
assert len(list(module.parameters(recurse=True))) == 2
optimizer = cpm.Optimizer(MyOptim(module.parameters()))
#--------------
# Iteration 1
#--------------
x = numpy.array([4], dtype)
### Forward
y = module(x)
assert isinstance(y, torch.Tensor)
numpy.testing.assert_array_equal(y.detach().numpy(), [24])
### Backward
y.backward()
numpy.testing.assert_array_equal(link.l1.p1.grad, [12])
numpy.testing.assert_array_equal(link.p2.grad, [8])
### Optimizer step
optimizer.set_constant(3)
optimizer.step()
# (Torch grads are only synchronized after step())
numpy.testing.assert_array_equal(module.l1.p1.grad.detach().numpy(), [12])
numpy.testing.assert_array_equal(module.p2.grad.detach().numpy(), [8])
numpy.testing.assert_array_equal(link.p2.array, [3])
numpy.testing.assert_array_equal(link.l1.p1.array, [3])
numpy.testing.assert_array_equal(module.p2.detach().numpy(), [3])
numpy.testing.assert_array_equal(module.l1.p1.detach().numpy(), [3])
### Zero grad
optimizer.zero_grad()
numpy.testing.assert_array_equal(link.l1.p1.grad, [0])
numpy.testing.assert_array_equal(link.p2.grad, [0])
numpy.testing.assert_array_equal(module.l1.p1.grad.detach().numpy(), [0])
numpy.testing.assert_array_equal(module.p2.grad.detach().numpy(), [0])
#--------------
# Iteration 2
#--------------
x = numpy.array([5], dtype)
### Forward
y = module(x)
assert isinstance(y, torch.Tensor)
numpy.testing.assert_array_equal(y.detach().numpy(), [45])
### Backward
y.backward()
numpy.testing.assert_array_equal(link.l1.p1.grad, [15])
numpy.testing.assert_array_equal(link.p2.grad, [15])
### Optimizer step
optimizer.set_constant(9)
optimizer.step()
# (Torch grads are only synchronized after step())
numpy.testing.assert_array_equal(module.l1.p1.grad.detach().numpy(), [15])
numpy.testing.assert_array_equal(module.p2.grad.detach().numpy(), [15])
numpy.testing.assert_array_equal(link.p2.array, [9])
numpy.testing.assert_array_equal(link.l1.p1.array, [9])
numpy.testing.assert_array_equal(module.p2.detach().numpy(), [9])
numpy.testing.assert_array_equal(module.l1.p1.detach().numpy(), [9])
### Zero grad
optimizer.zero_grad()
numpy.testing.assert_array_equal(link.l1.p1.grad, [0])
numpy.testing.assert_array_equal(link.p2.grad, [0])
numpy.testing.assert_array_equal(module.l1.p1.grad.detach().numpy(), [0])
numpy.testing.assert_array_equal(module.p2.grad.detach().numpy(), [0])