How to use the xenonpy.model.training.extension.TensorConverter function in xenonpy

def test_persist_1(data):
    model = deepcopy(data[0])
    trainer = Trainer(model=model, optimizer=Adam(lr=0.1), loss_func=MSELoss(), epochs=200)
    trainer.extend(TensorConverter(), Persist('model_dir'))
    trainer.fit(*data[1], *data[1])

    persist = trainer['persist']
    checker = persist._checker
    assert isinstance(persist, Persist)
    assert isinstance(checker.model, torch.nn.Module)
    assert isinstance(checker.describe, dict)
    assert isinstance(checker.files, list)
    assert set(checker.files) == {'model', 'init_state', 'model_structure', 'describe', 'training_info', 'final_state'}

    trainer = Trainer.load(checker)
    assert isinstance(trainer.training_info, pd.DataFrame)
    assert isinstance(trainer.model, torch.nn.Module)
    assert isinstance(trainer._training_info, list)
    assert trainer.optimizer is None
    assert trainer.lr_scheduler is None

model = _Net(n_feature=2, n_hidden=10, n_output=2)

    n_data = np.ones((100, 2))
    x0 = np.random.normal(2 * n_data, 1)
    y0 = np.zeros(100)
    x1 = np.random.normal(-2 * n_data, 1)
    y1 = np.ones(100)

    x = np.vstack((x0, x1))
    y = np.concatenate((y0, y1))
    s = np.arange(x.shape[0])
    np.random.shuffle(s)
    x, y = x[s], y[s]

    trainer = Trainer(model=model, optimizer=Adam(lr=0.1), loss_func=CrossEntropyLoss(), epochs=200)
    trainer.extend(TensorConverter(x_dtype=torch.float32, y_dtype=torch.long, argmax=True))
    trainer.fit(x, y)

    y_p, y_t = trainer.predict(x, y)
    assert y_p.shape == (200,)
    assert np.all(y_p == y_t)

    # trainer.reset()
    val_set = DataLoader(ArrayDataset(x, y, dtypes=(torch.float, torch.long)), batch_size=20)
    trainer.extend(TensorConverter(x_dtype=torch.float32, y_dtype=torch.long, auto_reshape=False))
    y_p, y_t = trainer.predict(dataset=val_set)
    assert y_p.shape == (200, 2)

    y_p = np.argmax(y_p, 1)
    assert np.all(y_p == y_t)

y = np.concatenate((y0, y1))
    s = np.arange(x.shape[0])
    np.random.shuffle(s)
    x, y = x[s], y[s]

    trainer = Trainer(model=model, optimizer=Adam(lr=0.1), loss_func=CrossEntropyLoss(), epochs=200)
    trainer.extend(TensorConverter(x_dtype=torch.float32, y_dtype=torch.long, argmax=True))
    trainer.fit(x, y)

    y_p, y_t = trainer.predict(x, y)
    assert y_p.shape == (200,)
    assert np.all(y_p == y_t)

    # trainer.reset()
    val_set = DataLoader(ArrayDataset(x, y, dtypes=(torch.float, torch.long)), batch_size=20)
    trainer.extend(TensorConverter(x_dtype=torch.float32, y_dtype=torch.long, auto_reshape=False))
    y_p, y_t = trainer.predict(dataset=val_set)
    assert y_p.shape == (200, 2)

    y_p = np.argmax(y_p, 1)
    assert np.all(y_p == y_t)

def test_trainer_prediction_1(data):
    model = deepcopy(data[0])
    trainer = Trainer(model=model, optimizer=Adam(lr=0.1), loss_func=MSELoss(), epochs=200)
    trainer.extend(TensorConverter())
    trainer.fit(*data[1], *data[1])

    trainer = Trainer(model=model).extend(TensorConverter())
    y_p = trainer.predict(data[1][0])
    assert np.any(np.not_equal(y_p, data[1][1].numpy()))
    assert np.allclose(y_p, data[1][1].numpy(), rtol=0, atol=0.2)

    y_p, y_t = trainer.predict(*data[1])
    assert np.any(np.not_equal(y_p, y_t))
    assert np.allclose(y_p, y_t, rtol=0, atol=0.2)

    val_set = DataLoader(TensorDataset(*data[1]), batch_size=50)
    y_p, y_t = trainer.predict(dataset=val_set)
    assert np.any(np.not_equal(y_p, y_t))
    assert np.allclose(y_p, y_t, rtol=0, atol=0.2)

# test auto reshape; #189
    converter = TensorConverter(auto_reshape=False)
    x, y = converter.input_proc(np_1, None, trainer=trainer)  # noqa
    assert isinstance(x, torch.Tensor)
    assert x.shape == (3,)

    x, y = converter.input_proc(se_1, None, trainer=trainer)  # noqa
    assert isinstance(x, torch.Tensor)
    assert x.shape == (3,)

    x, y = converter.input_proc(pd_1, None, trainer=trainer)  # noqa
    assert isinstance(x, torch.Tensor)
    assert x.shape == (3, 1)

    converter = TensorConverter()
    x, y = converter.input_proc(np_1, None, trainer=trainer)  # noqa
    assert isinstance(x, torch.Tensor)
    assert x.shape == (3, 1)

    x, y = converter.input_proc(se_1, None, trainer=trainer)  # noqa
    assert isinstance(x, torch.Tensor)
    assert x.shape == (3, 1)

    x, y = converter.input_proc(pd_1, None, trainer=trainer)  # noqa
    assert isinstance(x, torch.Tensor)
    assert x.shape == (3, 1)

    # normal tests
    x, y = converter.input_proc(np_, None, trainer=trainer)  # noqa
    assert isinstance(x, torch.Tensor)
    assert x.shape == (2, 3)

def test_tensor_converter_2():

    class _Trainer(BaseRunner):

        def __init__(self):
            super().__init__()
            self.non_blocking = False

        def predict(self, x_, y_):
            return x_, y_

    trainer = _Trainer()
    converter = TensorConverter()
    np_ = np.asarray([[1, 2, 3], [4, 5, 6]])
    pd_ = pd.DataFrame(np_)
    tensor_ = torch.Tensor(np_)  # noqa

    x, y = converter.input_proc(np_, np_[0], trainer=trainer)  # noqa
    assert isinstance(y, torch.Tensor)
    assert y.shape == (3, 1)
    assert torch.equal(y, tensor_[0].unsqueeze(-1))

    x, y = converter.input_proc(pd_, pd_.iloc[0], trainer=trainer)  # noqa
    assert isinstance(y, torch.Tensor)
    assert y.shape == (3, 1)
    assert torch.equal(y, tensor_[0].unsqueeze(-1))

    x, y = converter.input_proc(tensor_, tensor_[0], trainer=trainer)  # noqa
    assert isinstance(y, torch.Tensor)

def test_tensor_converter_3():
    converter = TensorConverter()
    np_ = np.asarray([[1, 2, 3], [4, 5, 6]])
    tensor_ = torch.from_numpy(np_)

    y, y_ = converter.output_proc(tensor_, None, training=True)
    assert y_ is None
    assert isinstance(y, torch.Tensor)
    assert y.shape == (2, 3)
    assert torch.equal(y, tensor_)

    y, y_ = converter.output_proc(tensor_, tensor_, training=True)
    assert isinstance(y, torch.Tensor)
    assert isinstance(y_, torch.Tensor)
    assert y.equal(y_)
    assert y.shape == (2, 3)
    assert torch.equal(y, tensor_)

self.non_blocking = False

        def predict(self, x_, y_):  # noqa
            return x_, y_

    trainer = _Trainer()
    arr_1 = [1, 2, 3]
    np_1 = np.asarray(arr_1)
    se_1 = pd.Series(arr_1)
    pd_1 = pd.DataFrame(arr_1)
    np_ = np.asarray([arr_1, arr_1])
    pd_ = pd.DataFrame(np_)
    tensor_ = torch.Tensor(np_)

    # test auto reshape; #189
    converter = TensorConverter(auto_reshape=False)
    x, y = converter.input_proc(np_1, None, trainer=trainer)  # noqa
    assert isinstance(x, torch.Tensor)
    assert x.shape == (3,)

    x, y = converter.input_proc(se_1, None, trainer=trainer)  # noqa
    assert isinstance(x, torch.Tensor)
    assert x.shape == (3,)

    x, y = converter.input_proc(pd_1, None, trainer=trainer)  # noqa
    assert isinstance(x, torch.Tensor)
    assert x.shape == (3, 1)

    converter = TensorConverter()
    x, y = converter.input_proc(np_1, None, trainer=trainer)  # noqa
    assert isinstance(x, torch.Tensor)
    assert x.shape == (3, 1)

assert isinstance(y[0], torch.Tensor)
    assert torch.equal(y[0], tensor_)

    y, y_ = converter.output_proc(tensor_, tensor_, training=False)
    assert isinstance(y, np.ndarray)
    assert isinstance(y_, np.ndarray)
    assert np.all(y == y_)
    assert y.shape == (2, 3)
    assert np.all(y == tensor_.numpy())

    y, _ = converter.output_proc((tensor_,), None, training=False)
    assert isinstance(y, tuple)
    assert isinstance(y[0], np.ndarray)
    assert np.all(y[0] == tensor_.numpy())

    converter = TensorConverter(argmax=True)
    y, y_ = converter.output_proc(tensor_, tensor_, training=False)
    assert isinstance(y, np.ndarray)
    assert isinstance(y_, np.ndarray)
    assert y.shape == (2,)
    assert y_.shape == (2, 3)
    assert np.all(y == np.argmax(np_, 1))

    y, y_ = converter.output_proc((tensor_, tensor_), None, training=False)
    assert isinstance(y, tuple)
    assert y_ is None
    assert y[0].shape == (2,)
    assert y[0].shape == y[1].shape
    assert np.all(y[0] == np.argmax(np_, 1))

assert torch.equal(y, tensor_)

    x, y = converter.input_proc(tensor_, tensor_, trainer=trainer)  # noqa
    assert isinstance(x, torch.Tensor)
    assert x.shape == (2, 3)
    assert torch.equal(x, tensor_)
    assert torch.equal(y, tensor_)

    converter = TensorConverter(x_dtype=torch.long)
    x, y = converter.input_proc((np_, np_), np_, trainer=trainer)  # noqa
    assert isinstance(x, tuple)
    assert len(x) == 2
    assert x[0].dtype == torch.long
    assert x[1].dtype == torch.long

    converter = TensorConverter(x_dtype=(torch.long, torch.float32), y_dtype=torch.long)
    x, y = converter.input_proc((np_, np_), np_, trainer=trainer)  # noqa
    assert isinstance(x, tuple)
    assert len(x) == 2
    assert x[0].dtype == torch.long
    assert x[1].dtype == torch.float32
    assert y.dtype == torch.long

    converter = TensorConverter(x_dtype=(torch.long, torch.float32))
    x, y = converter.input_proc((pd_, pd_), pd_, trainer=trainer)  # noqa
    assert isinstance(x, tuple)
    assert len(x) == 2
    assert x[0].dtype == torch.long
    assert x[1].dtype == torch.float32

    # for tensor input, dtype change will never be executed
    converter = TensorConverter(x_dtype=(torch.long, torch.long))