def _convert_to_tensor(
self,
data: pd.DataFrame,
feature_columns: Optional[Union[List[str], List[List[str]], List[int],
List[List[int]]]] = None,
dtypes: Optional[torch.dtype] = None,
unsqueeze: bool = True,
) -> torch.Tensor:
"""Handle conversion of data to tensor.
Same arguments as in ``convert_pandas_to_torch_tensor``."""
# TODO(amog): Add `_convert_numpy_to_torch_tensor to use based on input type.
# Reduce conversion cost if input is in Numpy
if isinstance(feature_columns, dict):
features_tensor = {
key: convert_pandas_to_torch_tensor(
data,
feature_columns[key],
dtypes[key] if isinstance(dtypes, dict) else dtypes,
unsqueeze=unsqueeze,
)
for key in feature_columns
}
else:
features_tensor = convert_pandas_to_torch_tensor(
data,
columns=feature_columns,
column_dtypes=dtypes,
unsqueeze=unsqueeze,
)
return features_tensor
def test_multi_input(self):
tensors = convert_pandas_to_torch_tensor(data_batch,
columns=[["A"], ["B"]])
assert len(tensors) == 2
for i in range(len(tensors)):
tensor = tensors[i]
assert tensor.size() == (len(data_batch), 1)
assert np.array_equal(
tensor.numpy(), data_batch[[data_batch.columns[i]]].to_numpy())
def test_single_tensor_dtype(self):
tensor = convert_pandas_to_torch_tensor(data_batch,
column_dtypes=torch.float)
assert tensor.size() == (len(data_batch), len(data_batch.columns))
assert tensor.dtype == torch.float
assert np.array_equal(tensor.numpy(), data_batch.to_numpy())
def test_single_tensor_columns(self):
tensor = convert_pandas_to_torch_tensor(data_batch, columns=["A"])
assert tensor.size() == (len(data_batch), len(data_batch.columns) - 1)
assert np.array_equal(tensor.numpy(), data_batch[["A"]].to_numpy())
def test_invalid_args(self):
with pytest.raises(TypeError):
convert_pandas_to_torch_tensor(
data_batch,
columns=["A", "B"],
column_dtypes=[torch.float, torch.float])