Implement map_batches to align with TorchArrow API #359
Conversation
facebook-github-bot
added
the
CLA Signed
|
@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
|
|
||
| def __iter__(self) -> Iterator[T_co]: | ||
| batch: List = [] | ||
| for d in self.datapipe: |
There was a problem hiding this comment.
nit: can this be done by .batch(batch_size).map(...)?
There was a problem hiding this comment.
hhhh, totally make sense to me. 😃
There was a problem hiding this comment.
After testing, .batch(batch_size).map() is not equal to .map_batches() when input_col is specified because map would treat input as a single data structure rather than a batch. Then, the input_col would be applied at batch level rather than data level.
Let's say our input is [(0,1), (2,3), (3,4), (5,6), (7,8), (9,10)] and batch size as 3. With input_col as 1, the inputs sent into fn are different between the two implementations.
.batch().map(): Inputs are(2, 3)and(7, 8).map_batches(): inputs are(1, 3, 4)and(6, 8, 10)
There was a problem hiding this comment.
I see. This is because we have data[idx] rather than batch[idx] within _apply_fn.
|
@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
|
@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
Per title