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