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