File size: 2,780 Bytes
2568013 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# Copyright (C) 2024-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# utilitary functions for DUSt3R
# --------------------------------------------------------
import numpy as np
import torch
def todevice(batch, device, callback=None, non_blocking=False):
"""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.
"""
if callback:
batch = callback(batch)
if isinstance(batch, dict):
return {k: todevice(v, device) for k, v in batch.items()}
if isinstance(batch, (tuple, list)):
return type(batch)(todevice(x, device) for x in batch)
x = batch
if device == "numpy":
if isinstance(x, torch.Tensor):
x = x.detach().cpu().numpy()
elif x is not None:
if isinstance(x, np.ndarray):
x = torch.from_numpy(x)
if torch.is_tensor(x):
x = x.to(device, non_blocking=non_blocking)
return x
to_device = todevice # alias
def to_numpy(x):
return todevice(x, "numpy")
def to_cpu(x):
return todevice(x, "cpu")
def to_cuda(x):
return todevice(x, "cuda")
def collate_with_cat(whatever, lists=False):
if isinstance(whatever, dict):
return {k: collate_with_cat(vals, lists=lists) for k, vals in whatever.items()}
elif isinstance(whatever, (tuple, list)):
if len(whatever) == 0:
return whatever
elem = whatever[0]
T = type(whatever)
if elem is None:
return None
if isinstance(elem, (bool, float, int, str)):
return whatever
if isinstance(elem, tuple):
return T(collate_with_cat(x, lists=lists) for x in zip(*whatever))
if isinstance(elem, dict):
return {
k: collate_with_cat([e[k] for e in whatever], lists=lists) for k in elem
}
if isinstance(elem, torch.Tensor):
return listify(whatever) if lists else torch.cat(whatever)
if isinstance(elem, np.ndarray):
return (
listify(whatever)
if lists
else torch.cat([torch.from_numpy(x) for x in whatever])
)
# otherwise, we just chain lists
return sum(whatever, T())
def listify(elems):
return [x for e in elems for x in e]
|