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)
| 21 | |
| 22 | |
| 23 | def todevice(batch, device, callback=None, non_blocking=False): |
| 24 | """Transfer some variables to another device (i.e. GPU, CPU:torch, CPU:numpy). |
| 25 | |
| 26 | batch: list, tuple, dict of tensors or other things |
| 27 | device: pytorch device or 'numpy' |
| 28 | callback: function that would be called on every sub-elements. |
| 29 | """ |
| 30 | if callback: |
| 31 | batch = callback(batch) |
| 32 | |
| 33 | if isinstance(batch, dict): |
| 34 | return {k: todevice(v, device) for k, v in batch.items()} |
| 35 | |
| 36 | if isinstance(batch, (tuple, list)): |
| 37 | return type(batch)(todevice(x, device) for x in batch) |
| 38 | |
| 39 | x = batch |
| 40 | if device == "numpy": |
| 41 | if isinstance(x, torch.Tensor): |
| 42 | x = x.detach().cpu().numpy() |
| 43 | elif x is not None: |
| 44 | if isinstance(x, np.ndarray): |
| 45 | x = torch.from_numpy(x) |
| 46 | if torch.is_tensor(x): |
| 47 | x = x.to(device, non_blocking=non_blocking) |
| 48 | return x |
| 49 | |
| 50 | |
| 51 | to_device = todevice # alias |