Transfer some variables to another device (i.e. GPU, CPU:torch, CPU:numpy). batch: list, tuple, dict of tensors or other things device: pytorch device or 'numpy' callback: function that would be called on every sub-elements.
(batch, device, callback=None, non_blocking=False)
| 7 | |
| 8 | |
| 9 | def todevice(batch, device, callback=None, non_blocking=False): |
| 10 | """Transfer some variables to another device (i.e. GPU, CPU:torch, CPU:numpy). |
| 11 | |
| 12 | batch: list, tuple, dict of tensors or other things |
| 13 | device: pytorch device or 'numpy' |
| 14 | callback: function that would be called on every sub-elements. |
| 15 | """ |
| 16 | if callback: |
| 17 | batch = callback(batch) |
| 18 | |
| 19 | if isinstance(batch, dict): |
| 20 | return {k: todevice(v, device) for k, v in batch.items()} |
| 21 | |
| 22 | if isinstance(batch, (tuple, list)): |
| 23 | return type(batch)(todevice(x, device) for x in batch) |
| 24 | |
| 25 | x = batch |
| 26 | if device == "numpy": |
| 27 | if isinstance(x, torch.Tensor): |
| 28 | x = x.detach().cpu().numpy() |
| 29 | elif x is not None: |
| 30 | if isinstance(x, np.ndarray): |
| 31 | x = torch.from_numpy(x) |
| 32 | if torch.is_tensor(x): |
| 33 | x = x.to(device, non_blocking=non_blocking) |
| 34 | return x |
| 35 | |
| 36 | |
| 37 | to_device = todevice # alias |