Spaces:
Running
Running
File size: 7,574 Bytes
5221c8c | 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 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 | import torch
import numpy as np
torch.manual_seed(0)
np.random.seed(0)
import GPUtil, gc
from pynvml import nvmlInit, nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo
use_cuda = True
FloatTensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor
LongTensor = torch.cuda.LongTensor if use_cuda else torch.LongTensor
IntTensor = torch.cuda.IntTensor if use_cuda else torch.IntTensor
ByteTensor = torch.cuda.ByteTensor if use_cuda else torch.ByteTensor
BoolTensor = torch.cuda.BoolTensor if use_cuda else torch.BoolTensor
Tensor = FloatTensor
def set_device(device):
# Handle torch.device objects.
if isinstance(device, torch.device):
device = str(device)
globals()["use_cuda"] = device != "cpu"
globals()["FloatTensor"] = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor
globals()["LongTensor"] = torch.cuda.LongTensor if use_cuda else torch.LongTensor
globals()["IntTensor"] = torch.cuda.IntTensor if use_cuda else torch.IntTensor
globals()["ByteTensor"] = torch.cuda.ByteTensor if use_cuda else torch.ByteTensor
globals()["BoolTensor"] = torch.cuda.BoolTensor if use_cuda else torch.BoolTensor
globals()["Tensor"] = FloatTensor
torch_device = torch.device(device if (torch.cuda.is_available()) else "cpu")
if torch.cuda.is_available() and device.startswith("cuda"):
torch.cuda.set_device(torch_device)
def print_gpu_usage(gpu_no=0):
GPUtil.showUtilization()
nvmlInit()
h = nvmlDeviceGetHandleByIndex(gpu_no)
info = nvmlDeviceGetMemoryInfo(h)
print(f"total : {info.total}")
print(f"free : {info.free}")
print(f"used : {info.used}")
print(torch.cuda.memory_summary())
for obj in gc.get_objects():
try:
if torch.is_tensor(obj) or (
hasattr(obj, "data") and torch.is_tensor(obj.data)
):
print(type(obj), obj.size())
except:
pass
# def proj_tensor(u, v):
# # u, v: [B, ..., D]
# # projet v to u
# B, D = u.shape[0], u.shape[-1]
# uv = torch.sum(u * v, axis=-1)
# uu = torch.sum(u * u, axis=-1)
# a = (uv / uu).unsqueeze(dim=-1)
# repeat_dim = [1] * len(u.shape)
# repeat_dim[-1] = D
# a = a.repeat(repeat_dim)
# return a * u
def tensor_q2qR(q):
"""
input q: [..., T, q_dim(=6)]
output qR: [..., T, 3, 3]
"""
q_shape = q.shape
q_reshape = tuple(list(q_shape[:-1]) + [2, 3])
q_ = q.reshape(q_reshape) # [..., T, 2, 3]
v1 = q_[..., 0, :]
v2 = q_[..., 1, :]
e1 = torch.nn.functional.normalize(v1, dim=-1)
u2 = v2 - (e1 * v2).sum(-1, keepdim=True) * e1 # v2 - proj_tensor(v1, v2)
e2 = torch.nn.functional.normalize(u2, dim=-1)
e3 = torch.cross(e1, e2, dim=-1)
return torch.stack((e1, e2, e3), dim=-1)
# a1, a2 = d6[..., :3], d6[..., 3:]
# b1 = torch.nn.functional.normalize(a1, dim=-1)
# b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1
# b2 = torch.nn.functional.normalize(b2, dim=-1)
# b3 = torch.cross(b1, b2, dim=-1)
# return torch.stack((b1, b2, b3), dim=-2)
from fairmotion.utils import constants
def tensor_r_to_rT(r_dn, apply_height=False):
"""
input r_dn : [..., r_dim]
output rT : [..., 4, 4]
"""
dtheta, dx, dz, h = r_dn[..., 0], r_dn[..., 1], r_dn[..., 2], r_dn[..., 3]
dcos, dsin = torch.cos(dtheta), torch.sin(dtheta)
repeat_shape = tuple(list(r_dn.shape[:-1]) + [1, 1])
root_T = Tensor(np.tile(constants.eye_T(), repeat_shape))
root_T[..., 0, 0] = dcos
root_T[..., 0, 2] = dsin
root_T[..., 0, 3] = dx
root_T[..., 2, 0] = -dsin
root_T[..., 2, 2] = dcos
root_T[..., 2, 3] = dz
if apply_height:
root_T[..., 1, 3] = h
return root_T
def tensor_p2T(p):
reshape = tuple(list(p.shape[:-1]) + [4, 4])
T = Tensor(constants.eye_T()).expand(*reshape).clone()
T[..., :3, 3] = p
return T
def cdn(torch_tensor):
return torch_tensor.cpu().detach().numpy()
# below are rotation_conversions code copied from pytorch3d
# https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html
def _copysign(a, b):
"""
Return a tensor where each element has the absolute value taken from the,
corresponding element of a, with sign taken from the corresponding
element of b. This is like the standard copysign floating-point operation,
but is not careful about negative 0 and NaN.
Args:
a: source tensor.
b: tensor whose signs will be used, of the same shape as a.
Returns:
Tensor of the same shape as a with the signs of b.
"""
signs_differ = (a < 0) != (b < 0)
return torch.where(signs_differ, -a, a)
def _sqrt_positive_part(x):
"""
Returns torch.sqrt(torch.max(0, x))
but with a zero subgradient where x is 0.
"""
ret = torch.zeros_like(x)
positive_mask = x > 0
ret[positive_mask] = torch.sqrt(x[positive_mask])
return ret
def matrix_to_quaternion(matrix):
"""
Convert rotations given as rotation matrices to quaternions.
Args:
matrix: Rotation matrices as tensor of shape (..., 3, 3).
Returns:
quaternions with real part first, as tensor of shape (..., 4).
"""
if matrix.size(-1) != 3 or matrix.size(-2) != 3:
raise ValueError(f"Invalid rotation matrix shape f{matrix.shape}.")
m00 = matrix[..., 0, 0]
m11 = matrix[..., 1, 1]
m22 = matrix[..., 2, 2]
o0 = 0.5 * _sqrt_positive_part(1 + m00 + m11 + m22)
x = 0.5 * _sqrt_positive_part(1 + m00 - m11 - m22)
y = 0.5 * _sqrt_positive_part(1 - m00 + m11 - m22)
z = 0.5 * _sqrt_positive_part(1 - m00 - m11 + m22)
o1 = _copysign(x, matrix[..., 2, 1] - matrix[..., 1, 2])
o2 = _copysign(y, matrix[..., 0, 2] - matrix[..., 2, 0])
o3 = _copysign(z, matrix[..., 1, 0] - matrix[..., 0, 1])
return torch.stack((o0, o1, o2, o3), -1)
def quaternion_to_axis_angle(quaternions):
"""
Convert rotations given as quaternions to axis/angle.
Args:
quaternions: quaternions with real part first,
as tensor of shape (..., 4).
Returns:
Rotations given as a vector in axis angle form, as a tensor
of shape (..., 3), where the magnitude is the angle
turned anticlockwise in radians around the vector's
direction.
"""
norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True)
half_angles = torch.atan2(norms, quaternions[..., :1])
angles = 2 * half_angles
eps = 1e-6
small_angles = angles.abs() < eps
sin_half_angles_over_angles = torch.empty_like(angles)
sin_half_angles_over_angles[~small_angles] = (
torch.sin(half_angles[~small_angles]) / angles[~small_angles]
)
# for x small, sin(x/2) is about x/2 - (x/2)^3/6
# so sin(x/2)/x is about 1/2 - (x*x)/48
sin_half_angles_over_angles[small_angles] = (
0.5 - (angles[small_angles] * angles[small_angles]) / 48
)
return quaternions[..., 1:] / sin_half_angles_over_angles
def matrix_to_axis_angle(matrix):
"""
Convert rotations given as rotation matrices to axis/angle.
Args:
matrix: Rotation matrices as tensor of shape (..., 3, 3).
Returns:
Rotations given as a vector in axis angle form, as a tensor
of shape (..., 3), where the magnitude is the angle
turned anticlockwise in radians around the vector's
direction.
"""
return quaternion_to_axis_angle(matrix_to_quaternion(matrix))
|