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