def __init__(self, height, width, lr = 1, aux_loss = False, ray_tracing = False):
super(Depth3DGridGen_with_mask, self).__init__()
self.height, self.width = height, width
self.aux_loss = aux_loss
self.lr = lr
self.ray_tracing = ray_tracing
self.grid = np.zeros( [self.height, self.width, 3], dtype=np.float32)
self.grid[:,:,0] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/self.height), 0), repeats = self.width, axis = 0).T, 0)
self.grid[:,:,1] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/self.width), 0), repeats = self.height, axis = 0), 0)
self.grid[:,:,2] = np.ones([self.height, width])
self.grid = torch.from_numpy(self.grid.astype(np.float32))
self.theta = self.grid[:,:,0] * np.pi/2 + np.pi/2
self.phi = self.grid[:,:,1] * np.pi
self.x = torch.sin(self.theta) * torch.cos(self.phi)
self.y = torch.sin(self.theta) * torch.sin(self.phi)
self.z = torch.cos(self.theta)
self.grid3d = torch.from_numpy(np.zeros( [self.height, self.width, 4], dtype=np.float32))
self.grid3d[:,:,0] = self.x
self.grid3d[:,:,1] = self.y
self.grid3d[:,:,2] = self.z
self.grid3d[:,:,3] = self.grid[:,:,2]
def forward(self, depth, trans0, trans1, rotate):
self.batchgrid3d = torch.zeros(torch.Size([depth.size(0)]) + self.grid3d.size())
for i in range(depth.size(0)):
self.batchgrid3d[i] = self.grid3d
self.batchgrid3d = Variable(self.batchgrid3d)
self.batchgrid = torch.zeros(torch.Size([depth.size(0)]) + self.grid.size())
for i in range(depth.size(0)):
self.batchgrid[i] = self.grid
self.batchgrid = Variable(self.batchgrid)
if depth.is_cuda:
self.batchgrid = self.batchgrid.cuda()
self.batchgrid3d = self.batchgrid3d.cuda()
x_ = self.batchgrid3d[:,:,:,0:1] * depth + trans0.view(-1,1,1,1).repeat(1, self.height, self.width, 1)
y_ = self.batchgrid3d[:,:,:,1:2] * depth + trans1.view(-1,1,1,1).repeat(1, self.height, self.width, 1)
z = self.batchgrid3d[:,:,:,2:3] * depth
#print(x.size(), y.size(), z.size())
rotate_z = rotate.view(-1,1,1,1).repeat(1,self.height, self.width,1) * np.pi
x = x_ * torch.cos(rotate_z) - y_ * torch.sin(rotate_z)
y = x_ * torch.sin(rotate_z) + y_ * torch.cos(rotate_z)
r = torch.sqrt(x**2 + y**2 + z**2) + 1e-5
#print(r)
theta = torch.acos(z/r)/(np.pi/2) - 1
#phi = torch.atan(y/x)
if depth.is_cuda:
phi = torch.atan(y/(x + 1e-5)) + np.pi * x.lt(0).type(torch.cuda.FloatTensor) * (y.ge(0).type(torch.cuda.FloatTensor) - y.lt(0).type(torch.cuda.FloatTensor))
else:
phi = torch.atan(y/(x + 1e-5)) + np.pi * x.lt(0).type(torch.FloatTensor) * (y.ge(0).type(torch.FloatTensor) - y.lt(0).type(torch.FloatTensor))
phi = phi/np.pi
output = torch.cat([theta,phi], 3)
return output
def forward(self, X, Z=None, diag=False):
if diag:
return self._diag(X)
variance = self.get_param("variance")
r = self._scaled_dist(X, Z)
return variance * torch.cos(r)
def draw(
self, n: int = 1, out: Optional[Tensor] = None, dtype: torch.dtype = torch.float
) -> Optional[Tensor]:
r"""Draw `n` qMC samples from the standard Normal.
Args:
n: The number of samples to draw.
out: An option output tensor. If provided, draws are put into this
tensor, and the function returns None.
dtype: The desired torch data type (ignored if `out` is provided).
Returns:
A `n x d` tensor of samples if `out=None` and `None` otherwise.
"""
# get base samples
samples = self._sobol_engine.draw(n, dtype=dtype)
if self._inv_transform:
# apply inverse transform (values to close to 0/1 result in inf values)
v = 0.5 + (1 - 1e-10) * (samples - 0.5)
samples_tf = torch.erfinv(2 * v - 1) * math.sqrt(2)
else:
# apply Box-Muller transform (note: [1] indexes starting from 1)
even = torch.arange(0, samples.shape[-1], 2)
Rs = (-2 * torch.log(samples[:, even])).sqrt()
thetas = 2 * math.pi * samples[:, 1 + even]
cos = torch.cos(thetas)
sin = torch.sin(thetas)
samples_tf = torch.stack([Rs * cos, Rs * sin], -1).reshape(n, -1)
# make sure we only return the number of dimension requested
samples_tf = samples_tf[:, : self._d]
if out is None:
return samples_tf
else:
out.copy_(samples_tf)
def neg_branin(X: Tensor) -> Tensor:
r"""Negative Branin test function.
Two-dimensional function (usually evaluated on `[-5, 10] x [0, 15]`):
`B(x) = (x2 - b x_1^2 + c x_1 - r)^2 + 10 (1-t) cos(x_1) + 10`
B has 3 minimizers for its global minimum at
`z_1 = (-pi, 12.275), z_2 = (pi, 2.275), z_3 = (9.42478, 2.475)`
with `B(z_i) = -0.397887`
Args:
X: A Tensor of size `2` or `k x 2` (`k` batch evaluations).
Returns:
`-B(X)`, the negative value of the standard Branin function.
"""
batch = X.ndimension() > 1
X = X if batch else X.unsqueeze(0)
t1 = X[:, 1] - 5.1 / (4 * math.pi ** 2) * X[:, 0] ** 2 + 5 / math.pi * X[:, 0] - 6
t2 = 10 * (1 - 1 / (8 * math.pi)) * torch.cos(X[:, 0])
B = t1 ** 2 + t2 + 10
result = -B
return result if batch else result.squeeze(0)
def _get_random_data(n, **tkwargs):
train_x1 = torch.linspace(0, 0.95, n + 1, **tkwargs) + 0.05 * torch.rand(
n + 1, **tkwargs
)
train_x2 = torch.linspace(0, 0.95, n, **tkwargs) + 0.05 * torch.rand(n, **tkwargs)
train_y1 = torch.sin(train_x1 * (2 * math.pi)) + 0.2 * torch.randn_like(train_x1)
train_y2 = torch.cos(train_x2 * (2 * math.pi)) + 0.2 * torch.randn_like(train_x2)
return train_x1.unsqueeze(-1), train_x2.unsqueeze(-1), train_y1, train_y2
def backward(ctx, grad_output):
input, = ctx.saved_tensors
grad_input = torch.stack((grad_output, torch.zeros_like(grad_output)), dim=len(grad_output.shape))
phase_input = angle(input)
phase_input = torch.stack((torch.cos(phase_input), torch.sin(phase_input)), dim=len(grad_output.shape))
grad_input = multiply_complex(phase_input, grad_input)
return 0.5*grad_input
def forward(ctx, input):
assert input.shape[-1]==2, "Complex tensor should have real and imaginary parts."
output = input.clone()
amplitude = torch.exp(input[..., 0])
# amplitude = input[..., 0]
output[..., 0] = amplitude*torch.cos(input[..., 1])
output[..., 1] = amplitude*torch.sin(input[..., 1])
ctx.save_for_backward(output)
return output
def __init__(self, input_dim: int, max_len: int = 5000) -> None:
super().__init__()
# Compute the positional encodings once in log space.
positional_encoding = torch.zeros(max_len, input_dim, requires_grad=False)
position = torch.arange(0, max_len).unsqueeze(1).float()
div_term = torch.exp(torch.arange(0, input_dim, 2).float() * -(math.log(10000.0) / input_dim))
positional_encoding[:, 0::2] = torch.sin(position * div_term)
positional_encoding[:, 1::2] = torch.cos(position * div_term)
positional_encoding = positional_encoding.unsqueeze(0)
self.register_buffer('positional_encoding', positional_encoding)
def __init__(self, d_model, dropout, max_len=5000):
super(PositionalEncoding, self).__init__()
self.dropout = torch.nn.Dropout(p=dropout)
pe = torch.zeros(max_len, d_model)
position = torch.arange(0., max_len).unsqueeze(1)
div_term = torch.exp(torch.arange(0., d_model, 2) * -(math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0)
self.register_buffer("pe", pe)
def __init__(self, dropout, dim, max_len=5000):
pe = torch.zeros(max_len, dim)
position = torch.arange(0, max_len).unsqueeze(1)
div_term = torch.exp(torch.arange(0, dim, 2) *
-(math.log(10000.0) / dim))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(1)
super(PositionalEncoding, self).__init__()
self.register_buffer('pe', pe)
self.dropout = nn.Dropout(p=dropout)
self.dim = dim
def _get_random_mt_data(**tkwargs):
train_x = torch.linspace(0, 0.95, 10, **tkwargs) + 0.05 * torch.rand(10, **tkwargs)
train_y1 = torch.sin(train_x * (2 * math.pi)) + torch.randn_like(train_x) * 0.2
train_y2 = torch.cos(train_x * (2 * math.pi)) + torch.randn_like(train_x) * 0.2
train_i_task1 = torch.full_like(train_x, dtype=torch.long, fill_value=0)
train_i_task2 = torch.full_like(train_x, dtype=torch.long, fill_value=1)
full_train_x = torch.cat([train_x, train_x])
full_train_i = torch.cat([train_i_task1, train_i_task2])
full_train_y = torch.cat([train_y1, train_y2])
train_X = torch.stack([full_train_x, full_train_i.type_as(full_train_x)], dim=-1)
train_Y = full_train_y
return train_X, train_Y
def add_positional_features(tensor: torch.Tensor,
min_timescale: float = 1.0,
max_timescale: float = 1.0e4):
# pylint: disable=line-too-long
"""
Implements the frequency-based positional encoding described
in `Attention is all you Need
<https://www.semanticscholar.org/paper/Attention-Is-All-You-Need-Vaswani-Shazeer/0737da0767d77606169cbf4187b83e1ab62f6077>`_ .
Adds sinusoids of different frequencies to a ``Tensor``. A sinusoid of a
different frequency and phase is added to each dimension of the input ``Tensor``.
This allows the attention heads to use absolute and relative positions.
The number of timescales is equal to hidden_dim / 2 within the range
(min_timescale, max_timescale). For each timescale, the two sinusoidal
signals sin(timestep / timescale) and cos(timestep / timescale) are
generated and concatenated along the hidden_dim dimension.
Parameters
----------
tensor : ``torch.Tensor``
a Tensor with shape (batch_size, timesteps, hidden_dim).
min_timescale : ``float``, optional (default = 1.0)
The smallest timescale to use.
max_timescale : ``float``, optional (default = 1.0e4)
The largest timescale to use.
Returns
-------
The input tensor augmented with the sinusoidal frequencies.
"""
_, timesteps, hidden_dim = tensor.size()
timestep_range = get_range_vector(timesteps, get_device_of(tensor)).data.float()
# We're generating both cos and sin frequencies,
# so half for each.
num_timescales = hidden_dim // 2
timescale_range = get_range_vector(num_timescales, get_device_of(tensor)).data.float()
log_timescale_increments = math.log(float(max_timescale) / float(min_timescale)) / float(num_timescales - 1)
inverse_timescales = min_timescale * torch.exp(timescale_range * -log_timescale_increments)
# Broadcasted multiplication - shape (timesteps, num_timescales)
scaled_time = timestep_range.unsqueeze(1) * inverse_timescales.unsqueeze(0)
# shape (timesteps, 2 * num_timescales)
sinusoids = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], 1)
if hidden_dim % 2 != 0:
# if the number of dimensions is odd, the cos and sin
# timescales had size (hidden_dim - 1) / 2, so we need
# to add a row of zeros to make up the difference.
sinusoids = torch.cat([sinusoids, sinusoids.new_zeros(timesteps, 1)], 1)
return tensor + sinusoids.unsqueeze(0)
def invert(self, head_coords, segment_len, mean_angle, eigenworms):
angles = torch.matmul(eigenworms, self.eigen_components)
angles += mean_angle.view(-1, 1)
ske_x = torch.sin(angles).view(-1, self.n_angles, 1)
ske_y = torch.cos(angles).view(-1, self.n_angles, 1)
skels_n = torch.cat([ske_x, ske_y], 2)*segment_len.view(-1, 1, 1)
skels_n = torch.cat([head_coords.view(-1, 1, 2), skels_n], 1)
skels_n = torch.cumsum(skels_n, dim=1)
return skels_n
def __init__(self, _, **kwargs):
super(PositionalLookupTableEmbeddings, self).__init__(_, **kwargs)
self.dropout = nn.Dropout(kwargs.get('dropout', 0.1))
# This could get us in trouble, if in doubt, pick something big
mxlen = kwargs.get('mxlen', 1000)
max_timescale = kwargs.get('max_timescale', 1.0e4)
log_timescale_increment = math.log(max_timescale) / self.dsz
inv_timescales = torch.exp(torch.arange(0, self.dsz, 2).float() * -log_timescale_increment)
pe = torch.zeros(mxlen, self.dsz)
position = torch.arange(0, mxlen).float().unsqueeze(1)
pe[:, 0::2] = torch.sin(position * inv_timescales)
pe[:, 1::2] = torch.cos(position * inv_timescales)
pe = pe.unsqueeze(0)
self.register_buffer('pe', pe)
def positional_encodings_like(x, t=None):
if t is None:
positions = torch.arange(0, x.size(1))
if x.is_cuda:
positions = positions.cuda(x.get_device())
else:
positions = t
encodings = torch.zeros(*x.size()[1:])
if x.is_cuda:
encodings = encodings.cuda(x.get_device())
for channel in range(x.size(-1)):
if channel % 2 == 0:
encodings[:, channel] = torch.sin(
positions / 10000 ** (channel / x.size(2)))
else:
encodings[:, channel] = torch.cos(
positions / 10000 ** ((channel - 1) / x.size(2)))
return Variable(encodings)
def _h_eigenworms_inv_T(head_x, head_y, segment_l,
mean_angle, eigenworms):
'''
Convert the eigen value transformed data into xy coordinates
'''
n_components = eigenworms.size(0)
angles = torch.mm(eigenworms.view(1, -1), EIGENWORMS_COMPONENTS_T[:n_components])
angles += mean_angle
ske_x = torch.sin(angles)*segment_l
ske_x = torch.cat([head_x.view(1,1), ske_x], 1)
ske_x = torch.cumsum(ske_x, dim=1)
ske_y = torch.cos(angles)*segment_l
ske_y = torch.cat([head_y.view(1,1), ske_y], 1)
ske_y = torch.cumsum(ske_y, dim=1)
skels_n = torch.cat((ske_x.view(-1, 1), ske_y.view(-1, 1)), 1)
return skels_n
def warmup_cosine(x, warmup=0.002):
s = 1 if x <= warmup else 0
return s*(x/warmup) + (1-s)*(0.5 * (1 + torch.cos(math.pi * x)))
from torch.autograd import Variable
from torch.nn import Parameter
from gpytorch.parameters import MLEParameterGroup
from gpytorch.kernels import RBFKernel, IndexKernel
from gpytorch.means import ConstantMean
from gpytorch.likelihoods import GaussianLikelihood
from gpytorch.inference import Inference
from gpytorch.random_variables import GaussianRandomVariable
# Simple training data: let's try to learn a sine function
train_x = Variable(torch.linspace(0, 1, 11))
y1_inds = Variable(torch.zeros(11).long())
y2_inds = Variable(torch.ones(11).long())
train_y1 = Variable(torch.sin(train_x.data * (2 * math.pi)))
train_y2 = Variable(torch.cos(train_x.data * (2 * math.pi)))
test_x = Variable(torch.linspace(0, 1, 51))
y1_inds_test = Variable(torch.zeros(51).long())
y2_inds_test = Variable(torch.ones(51).long())
test_y1 = Variable(torch.sin(test_x.data * (2 * math.pi)))
test_y2 = Variable(torch.cos(test_x.data * (2 * math.pi)))
class MultitaskGPModel(gpytorch.ObservationModel):
def __init__(self):
super(MultitaskGPModel, self).__init__(GaussianLikelihood())
self.mean_module = ConstantMean()
self.covar_module = RBFKernel()
self.task_covar_module = IndexKernel()
self.model_params = MLEParameterGroup(
def project(self, X, frame='w'):
"""
Projects 3D points onto the image plane
Parameters
----------
X : torch.Tensor [B,3,H,W]
3D points to be projected
frame : 'w'
Reference frame: 'c' for camera and 'w' for world
Returns
-------
points : torch.Tensor [B,H,W,2]
2D projected points that are within the image boundaries
"""
B, C, H, W = X.shape
assert C == 3
# World to camera:
if frame == 'c':
Xc = X.view(B, 3, -1) # [B, 3, HW]
elif frame == 'w':
Xc = (self.Tcw @ X).view(B, 3, -1) # [B, 3, HW]
else:
raise ValueError('Unknown reference frame {}'.format(frame))
c1 = self.poly_coeffs[:, 0].unsqueeze(1)
c2 = self.poly_coeffs[:, 1].unsqueeze(1)
c3 = self.poly_coeffs[:, 2].unsqueeze(1)
c4 = self.poly_coeffs[:, 3].unsqueeze(1)
# Project 3D points onto the camera image plane
X = Xc[:, 0] # [B, HW]
Y = Xc[:, 1] # [B, HW]
Z = Xc[:, 2] # [B, HW]
phi = torch.atan2(Y, X) # [B, HW]
rc = torch.sqrt(torch.pow(X, 2) + torch.pow(Y, 2)) # [B, HW]
theta_1 = math.pi / 2 - torch.atan2(Z, rc) # [B, HW]
theta_2 = torch.pow(theta_1, 2) # [B, HW]
theta_3 = torch.pow(theta_1, 3) # [B, HW]
theta_4 = torch.pow(theta_1, 4) # [B, HW]
rho = c1 * theta_1 + c2 * theta_2 + c3 * theta_3 + c4 * theta_4 # [B, HW]
rho = rho * ((X != 0) | (Y != 0) | (Z != 0))
u = rho * torch.cos(phi) / self.scale_factors[:, 0].unsqueeze(1) + self.principal_point[:, 0].unsqueeze(1) # [B, HW]
v = rho * torch.sin(phi) / self.scale_factors[:, 1].unsqueeze(1) + self.principal_point[:, 1].unsqueeze(1) # [B, HW]
# Normalize points
Xnorm = 2 * u / (W - 1)# - 1.
Ynorm = 2 * v / (H - 1)# - 1.
# Clamp out-of-bounds pixels
Xmask = ((Xnorm > 1) + (Xnorm < -1)).detach()
Xnorm[Xmask] = 2.
Ymask = ((Ynorm > 1) + (Ynorm < -1)).detach()
Ynorm[Ymask] = 2.
mask = (theta_1 * 180 * 2 / np.pi > 190.0).detach()
Xnorm[mask] = 2.
Ynorm[mask] = 2.
# Return pixel coordinates
return torch.stack([Xnorm, Ynorm], dim=-1).view(B, H, W, 2)
def encode(q):
c = torch.cos(q)
s = torch.sin(q)
return torch.cat((c, s), dim=2)
def y_exact(self, t):
return -0.5 * t**4 * torch.cos(2 * t) + 0.5 * t**3 * torch.sin(
2 * t) + 0.25 * t**2 * torch.cos(
2 * t) - t**3 + 2 * t**4 + (math.pi - 0.25) * t**2
def forward(self, output, labels, seq_len): #trans_params = (R00,R01,R10,R11,Tx,Ty,Tz)
lossMSE = nn.MSELoss()
lossHuber = nn.SmoothL1Loss()
lossCrossEntropy = nn.CrossEntropyLoss()
lossBinEntropy = nn.BCELoss()
batch_size = labels.shape[0]
score_seq, loc_seq, box_seq = output
score_seq = score_seq.permute(1,0) #(B,S,1)
loc_seq = loc_seq.permute(1,0,2) #(B,S,7)
box_seq = box_seq.permute(1,0,2) #(B,S,7)
total_loss = 0
for B in range(batch_size):
target = labels[B,:seq_len,8] #(3,1)
car_loc = torch.squeeze(target.nonzero())
num_car = car_loc.nelement()
if(num_car!=0):
trans_params_hat = torch.zeros(seq_len,5).cuda()
rot = labels[B,car_loc, 3] # CHANGED removed negative sign
trans_params_hat[:num_car, 0] = torch.cos(rot)
trans_params_hat[:num_car, 1] = torch.sin(rot)
trans_params_hat[:num_car,2:5] = labels[B, car_loc, :3]
# trans_params_hat[:num_car,5:] = labels[B, car_loc, 5:8]
if(num_car!=seq_len):
zero_rot = torch.tensor(0).float()
trans_params_hat[num_car:, 0] = torch.cos(zero_rot)
trans_params_hat[num_car:, 1] = torch.sin(zero_rot)
trans_params_hat[num_car:,2:5] = torch.tensor([0,0,-4]).cuda()
# trans_params_hat[num_car:,5:] = torch.tensor([2,2.5,5]).cuda()
# if(num_car == 3):
# print("yes", B, trans_params_hat)
ind = hungarian_matching(loc_seq[B,:,2:5].detach().cpu(), labels[B,car_loc,:3].view(num_car,-1).detach().cpu())
ind = np.concatenate((ind,[i for i in range(seq_len) if i not in ind]))
# loc_seq_matched = loc_seq[B,ind].view(num_car,-1)
loc_seq_matched = loc_seq[B,ind]
score_seq_matched = score_seq[B, ind]
box_seq_matched = box_seq[B, ind, :5]
size_seq_matched = box_seq[B, ind, 5:]
# pos_ind = score_seq_matched > 0.5
# pos_ind = torch.squeeze(pos_ind.nonzero())
# num_pos_ind = pos_ind.nelement()
# pdb.set_trace()
trans_mat_1 = torch.eye(3).view(1,-1).repeat(seq_len,1).cuda()
trans_mat_1[:,0] = loc_seq_matched[:,0] # c
trans_mat_1[:,1] = -loc_seq_matched[:,1] # -s
trans_mat_1[:,3] = loc_seq_matched[:,1] # s
trans_mat_1[:,4] = loc_seq_matched[:,0] # c
trans_mat_1[:,2] = loc_seq_matched[:,2] #tx
trans_mat_1[:,5] = loc_seq_matched[:,3] #ty
trans_mat_1 = trans_mat_1.view(seq_len,3,3)
# trans_mat_1 = trans_mat_1.detach() # here we detach bcoz we aim to train two network seperately!
trans_mat_2 = torch.eye(3).view(1,-1).repeat(seq_len,1).cuda()
trans_mat_2[:,0] = box_seq_matched[:,0] # c
trans_mat_2[:,1] = -box_seq_matched[:,1] # -s
trans_mat_2[:,3] = box_seq_matched[:,1] # s
trans_mat_2[:,4] = box_seq_matched[:,0] # c
trans_mat_2[:,2] = box_seq_matched[:,2] #tx
trans_mat_2[:,5] = box_seq_matched[:,3] #ty
trans_mat_2 = trans_mat_2.view(seq_len,3,3)
resultant_trans = torch.bmm(trans_mat_1, trans_mat_2)
resultant_trans = resultant_trans.view(seq_len, 9)
final_trans_params = resultant_trans[:,[0,3,2,5]]
z = (loc_seq_matched[:,4] + box_seq_matched[:,4]).view(seq_len,-1)
final_trans_params = torch.cat((final_trans_params,z),1)
rotx = loc_seq_matched[:,0]**2 + loc_seq_matched[:,1]**2
reg_loss_1 = lossMSE(rotx, torch.tensor([1]).cuda().float())
rotx = final_trans_params[:,0]**2 + final_trans_params[:,1]**2
reg_loss_2 = lossMSE(rotx, torch.tensor([1]).cuda().float())
# reg_loss = reg_loss.mean()
loss_where = lossHuber(loc_seq_matched , trans_params_hat)
loss_what = lossBinEntropy(score_seq_matched, target)
loss_residual = lossHuber(final_trans_params[car_loc].view(num_car,-1) , trans_params_hat[car_loc].view(num_car,-1))
loss_size = lossHuber(size_seq_matched[car_loc].view(num_car,-1), labels[B,:num_car,5:8].view(num_car,-1))
total_seq_loss = loss_where + loss_what + 0.01*(reg_loss_1+ reg_loss_2) + 1.5*loss_residual + 0.5*loss_size
else:
loss_what = lossBinEntropy(score_seq[B], target)
total_seq_loss = loss_what
total_loss += total_seq_loss
return total_loss/batch_size
def rotate2d(x, y, ang):
if isinstance(x, t.Tensor):
(c, s) = (t.cos(ang), t.sin(ang))
else:
(c, s) = (math.cos(ang), math.sin(ang))
return (c * x - s * y, s * x + c * y)
def get_normals(points3d):
""" Computes normal maps from 3D cloud
:param points3d: Ordered 3D cloud with shape [batch_size x height x width x 3]
:return: Normal map with shape [batch_size x height x width x 3]
"""
batch_size = points3d.shape[0]
dtype = points3d.dtype
device = points3d.device
pi = torch.tensor([np.pi], dtype=dtype, device=device)
nan = torch.tensor([np.nan], dtype=dtype, device=device)
# Create a huge tensor to keep the points, as well as the points shifted one step in all directions
hugeTensor = torch.zeros(batch_size,
points3d.shape[1] + 2,
points3d.shape[2] + 2,
points3d.shape[3],
9,
dtype=points3d.dtype,
device=device)
hugeTensor[:, 1:points3d.shape[1] + 1, 1:points3d.shape[2] + 1, 0:3,
0] = points3d # original points in center
# Now add them in this order: top-left, top-center, top-right
hugeTensor[:, 0:points3d.shape[1], 0:points3d.shape[2], 0:3, 1] = points3d
hugeTensor[:, 0:points3d.shape[1], 1:points3d.shape[2] + 1, 0:3,
2] = points3d
hugeTensor[:, 0:points3d.shape[1], 2:points3d.shape[2] + 2, 0:3,
3] = points3d
# Now center-left, center-right
hugeTensor[:, 1:points3d.shape[1] + 1, 0:points3d.shape[2], 0:3,
4] = points3d
hugeTensor[:, 1:points3d.shape[1] + 1, 2:points3d.shape[2] + 2, 0:3,
5] = points3d
# Now bottom-left, bottom-cente1, bottom-right
hugeTensor[:, 2:points3d.shape[1] + 2, 0:points3d.shape[2], 0:3,
6] = points3d
hugeTensor[:, 2:points3d.shape[1] + 2, 1:points3d.shape[2] + 1, 0:3,
7] = points3d
hugeTensor[:, 2:points3d.shape[1] + 2, 2:points3d.shape[2] + 2, 0:3,
8] = points3d
# Done! Now compute the mean vector for each pixel.
meanPoints = hugeTensor.mean(4)
h = hugeTensor.shape[1]
w = hugeTensor.shape[2]
S = torch.zeros(batch_size, h, w, 3, 3, dtype=dtype, device=device)
for k in range(9):
hugeTensor[:, :, :, :, k] = hugeTensor[:, :, :, :, k] - meanPoints
S = S + hugeTensor[:, :, :, :, k].reshape(
batch_size, h, w, 1, 3) * hugeTensor[:, :, :, :, k].reshape(
batch_size, h, w, 3, 1)
S = S / 8
eye3 = torch.eye(3, dtype=dtype, device=device).unsqueeze(0)
# Now, compute the smallest eigenvector of each covariance matrix
mat = torch.tensor(np.array([[0, 1, 1], [0, 0, 1], [0, 0, 0]]),
dtype=dtype,
device=device)
mat = mat.unsqueeze(0)
p1 = ((S * mat)**2).sum(4).sum(3)
q = (S * eye3).sum(4).sum(3) / 3
p2 = ((((S - q.reshape(batch_size, q.shape[1], q.shape[2], 1, 1)) * eye3) *
eye3)**2).sum(4).sum(3)
p2 = p2 + 2 * p1
val_mask = p2 > 0
p2[~val_mask] = 1
p2 = torch.sqrt(p2 / 6.0)
# No more division by zero
B = (S - q.reshape(batch_size, q.shape[1], q.shape[2], 1, 1) *
eye3) / p2.reshape(batch_size, p2.shape[1], p2.shape[2], 1, 1)
detB = B[:, :, :, 0, 0] * B[:, :, :, 1, 1] * B[:, :, :, 2, 2] + \
B[:, :, :, 0, 1] * B[:, :, :, 1, 2] * B[:, :, :, 2, 0] + \
B[:, :, :, 0, 2] * B[:, :, :, 1, 0] * B[:, :, :, 2, 1] - \
B[:, :, :, 0, 2] * B[:, :, :, 1, 1] * B[:, :, :, 2, 0] - \
B[:, :, :, 0, 1] * B[:, :, :, 1, 0] * B[:, :, :, 2, 2] - \
B[:, :, :, 0, 0] * B[:, :, :, 1, 2] * B[:, :, :, 2, 1]
detB = detB / 2.0
detB = detB.clamp(min=-1, max=1)
lambda0 = 2.0 * p2 * torch.cos(torch.acos(detB) / 3.0 + 2 * pi / 3.0) + q
detM = (S[:, :, :, 0, 0] - lambda0) * (
S[:, :, :, 1, 1] - lambda0) - S[:, :, :, 0, 1] * S[:, :, :, 1, 0]
normals = torch.ones(batch_size,
lambda0.shape[1],
lambda0.shape[2],
3,
dtype=dtype,
device=device)
val_mask = val_mask & (detM != 0)
detM[~val_mask] = 1
normals[:, :, :, 0] = (S[:, :, :, 0, 2] * (lambda0 - S[:, :, :, 1, 1]) +
S[:, :, :, 1, 2] * S[:, :, :, 0, 1]) / detM
normals[:, :, :, 1] = (S[:, :, :, 1, 2] * (lambda0 - S[:, :, :, 0, 0]) +
S[:, :, :, 0, 2] * S[:, :, :, 1, 0]) / detM
normals_length = torch.sqrt((normals**2).sum(3))
normals = normals / normals_length.reshape(batch_size, normals.shape[1],
normals.shape[2], 1)
neg_offsets = ((normals * meanPoints).sum(3)).reshape(
batch_size, normals.shape[1], normals.shape[2], 1)
normals = normals * torch.sign(neg_offsets)
normals[~val_mask] = nan
lambda0[~val_mask] = nan
normals = normals[:, 1:normals.shape[1] - 1, 1:normals.shape[2] - 1, :]
return normals, lambda0
def test_cos(x, y):
c = torch.cos(torch.add(x, y))
return c
def warmup_cosine(x, warmup=0.002):
s = 1 if x <= warmup else 0
return s * (x / warmup) + (1 - s) * (0.5 * (1 + torch.cos(math.pi * x)))
t.set_num_threads(int) #设定用于并行化CPU操作的OpenMP线程数
#数学操作Math operations
input = t.Tensor([[1, 2][3, 4]])
value = 3
t.abs(input, out=None) #返回张量：取绝对值
t.acos(input, out=None) #返回张量：反余弦
t.add(input, value, out=None) #返回张量：加value值
#t.addcdiv(input, value=1, tensor1, tensor2, out=None) #用tensor2对tensor1逐元素相除，然后乘以标量值value 并加到tensor
#t.addcmul(input, value=1, tensor1, tensor2, out=None) #用tensor2对tensor1逐元素相乘，并对结果乘以标量值value然后加到tensor。 张量的形状不需要匹配，但元素数量必须一致。 如果输入是FloatTensor or DoubleTensor类型，则value 必须为实数，否则须为整数。
t.asin(input, out=None) #取反正弦
t.atan(input, out=None) #取反正切
#t.atan2(input1, input2, out=None) #返回一个新张量，包含两个输入张量input1和input2的反正切函数
t.ceil(input, out=None) #天井函数，对输入input张量每个元素向上取整, 即取不小于每个元素的最小整数，并返回结果到输出。
#t.clamp(input, min, max, out=None) #将输入input张量每个元素的夹紧到区间 [min,max]，并返回结果到一个新张量。
t.cos(input, out=None)
t.cosh(input, out=None)
t.div(input, value, out=None) #将input逐元素除以标量值value，并返回结果到输出张量out
t.exp(input, out=None)
t.floor(input, out=None) #床函数: 返回一个新张量，包含输入input张量每个元素的floor，即不小于元素的最大整数。
#t.fmod(input, divisor, out=None) #计算除法余数。 除数与被除数可能同时含有整数和浮点数。此时，余数的正负与被除数相同。
t.frac(input, out=None) #返回小数部分
#t.lerp(start, end, weight, out=None) #对两个张量以start，end做线性插值， 将结果返回到输出张量。
t.log(input, out=None)
t.log1p(input, out=None) #计算 input+1的自然对数 yi=log(xi+1)
t.mul(input, value,
out=None) #用标量值value乘以输入input的每个元素，并返回一个新的结果张量。 out=tensor∗value
t.neg(input, out=None) #返回一个新张量，包含输入input 张量按元素取负。 即， out=−1∗input
#t.pow(input, exponent, out=None) #对输入input的按元素求exponent次幂值，并返回结果张量。 幂值exponent 可以为单一 float 数或者与input相同元素数的张量。
t.reciprocal(input, out=None) #返回一个新张量，包含输入input张量每个元素的倒数，即 1.0/x。
#t.remainder(input, divisor, out=None) #返回一个新张量，包含输入input张量每个元素的除法余数。 除数与被除数可能同时包含整数或浮点数。余数与除数有相同的符号。
def test_pi_quaternion(self, device):
one = torch.tensor(1.).to(device)
quaternion_log = torch.tensor([1., 0., 0.]).to(device)
expected = torch.tensor([torch.sin(one), 0., 0.,
torch.cos(one)]).to(device)
assert_allclose(kornia.quaternion_log_to_exp(quaternion_log), expected)
def forward(self, input):
input = input.reshape(len(input), self.input_dim, -1) # [N, F, T]
input = input.permute(0, 2, 1) # [N, T, F]
time_step = input.shape[1]
for ts in range(time_step):
x = input[:, ts, :]
if len(self.states) == 0: # hasn't initialized yet
self.init_states(x)
self.get_constants(x)
p_tm1 = self.states[0]
h_tm1 = self.states[1]
S_re_tm1 = self.states[2]
S_im_tm1 = self.states[3]
time_tm1 = self.states[4]
B_U = self.states[5]
B_W = self.states[6]
frequency = self.states[7]
x_i = torch.matmul(x * B_W[0], self.W_i) + self.b_i
x_ste = torch.matmul(x * B_W[0], self.W_ste) + self.b_ste
x_fre = torch.matmul(x * B_W[0], self.W_fre) + self.b_fre
x_c = torch.matmul(x * B_W[0], self.W_c) + self.b_c
x_o = torch.matmul(x * B_W[0], self.W_o) + self.b_o
i = self.inner_activation(x_i +
torch.matmul(h_tm1 * B_U[0], self.U_i))
ste = self.inner_activation(x_ste +
torch.matmul(h_tm1 *
B_U[0], self.U_ste))
fre = self.inner_activation(x_fre +
torch.matmul(h_tm1 *
B_U[0], self.U_fre))
ste = torch.reshape(ste, (-1, self.hidden_dim, 1))
fre = torch.reshape(fre, (-1, 1, self.freq_dim))
f = ste * fre
c = i * self.activation(x_c +
torch.matmul(h_tm1 * B_U[0], self.U_c))
time = time_tm1 + 1
omega = torch.tensor(2 * np.pi) * time * frequency
re = torch.cos(omega)
im = torch.sin(omega)
c = torch.reshape(c, (-1, self.hidden_dim, 1))
S_re = f * S_re_tm1 + c * re
S_im = f * S_im_tm1 + c * im
A = torch.square(S_re) + torch.square(S_im)
A = torch.reshape(A, (-1, self.freq_dim)).float()
A_a = torch.matmul(A * B_U[0], self.U_a)
A_a = torch.reshape(A_a, (-1, self.hidden_dim))
a = self.activation(A_a + self.b_a)
o = self.inner_activation(x_o +
torch.matmul(h_tm1 * B_U[0], self.U_o))
h = o * a
p = torch.matmul(h, self.W_p) + self.b_p
self.states = [p, h, S_re, S_im, time, None, None, None]
self.states = []
return self.fc_out(p).squeeze()
def log_prob(self, value):
log_prob = self.concentration * torch.cos(value - self.loc)
log_prob = log_prob - math.log(2 * math.pi) - _log_modified_bessel_fn_0(self.concentration)
return log_prob
def evaluate_torsion(r12, r23, r34, torsion_params, explicit_forces=True):
# Calculate dihedral angles from vectors
crossA = torch.cross(r12, r23, dim=1)
crossB = torch.cross(r23, r34, dim=1)
crossC = torch.cross(r23, crossA, dim=1)
normA = torch.norm(crossA, dim=1)
normB = torch.norm(crossB, dim=1)
normC = torch.norm(crossC, dim=1)
normcrossB = crossB / normB.unsqueeze(1)
cosPhi = torch.sum(crossA * normcrossB, dim=1) / normA
sinPhi = torch.sum(crossC * normcrossB, dim=1) / normC
phi = -torch.atan2(sinPhi, cosPhi)
ntorsions = len(torsion_params[0]["idx"])
pot = torch.zeros(ntorsions,
dtype=r12.dtype,
layout=r12.layout,
device=r12.device)
if explicit_forces:
coeff = torch.zeros(ntorsions,
dtype=r12.dtype,
layout=r12.layout,
device=r12.device)
for i in range(0, len(torsion_params)):
idx = torsion_params[i]["idx"]
k0 = torsion_params[i]["params"][:, 0]
phi0 = torsion_params[i]["params"][:, 1]
per = torsion_params[i]["params"][:, 2]
if torch.all(per > 0): # AMBER torsions
angleDiff = per * phi[idx] - phi0
pot.scatter_add_(0, idx, k0 * (1 + torch.cos(angleDiff)))
if explicit_forces:
coeff.scatter_add_(0, idx, -per * k0 * torch.sin(angleDiff))
else: # CHARMM torsions
angleDiff = phi[idx] - phi0
angleDiff[angleDiff < -pi] = angleDiff[angleDiff < -pi] + 2 * pi
angleDiff[angleDiff > pi] = angleDiff[angleDiff > pi] - 2 * pi
pot.scatter_add_(0, idx, k0 * angleDiff**2)
if explicit_forces:
coeff.scatter_add_(0, idx, 2 * k0 * angleDiff)
# coeff.unsqueeze_(1)
force0, force1, force2, force3 = None, None, None, None
if explicit_forces:
# Taken from OpenMM
normDelta2 = torch.norm(r23, dim=1)
norm2Delta2 = normDelta2**2
forceFactor0 = (-coeff * normDelta2) / (normA**2)
forceFactor1 = torch.sum(r12 * r23, dim=1) / norm2Delta2
forceFactor2 = torch.sum(r34 * r23, dim=1) / norm2Delta2
forceFactor3 = (coeff * normDelta2) / (normB**2)
force0vec = forceFactor0.unsqueeze(1) * crossA
force3vec = forceFactor3.unsqueeze(1) * crossB
s = (forceFactor1.unsqueeze(1) * force0vec -
forceFactor2.unsqueeze(1) * force3vec)
force0 = -force0vec
force1 = force0vec + s
force2 = force3vec - s
force3 = -force3vec
return pot, (force0, force1, force2, force3)
def backward(ctx, grad_output):
(x, ) = ctx.saved_tensors
ret = torch.cos(x) / x - torch.sin(x) / (x * x)
ret[x.abs() < 1e-10] = 0.0
return ret * grad_output
def convert(box: _RawBoxType, from_mode: "BoxMode",
to_mode: "BoxMode") -> _RawBoxType:
"""
Args:
box: can be a k-tuple, k-list or an Nxk array/tensor, where k = 4 or 5
from_mode, to_mode (BoxMode)
Returns:
The converted box of the same type.
"""
if from_mode == to_mode:
return box
original_type = type(box)
is_numpy = isinstance(box, np.ndarray)
single_box = isinstance(box, (list, tuple))
if single_box:
assert len(box) == 4 or len(box) == 5, (
"BoxMode.convert takes either a k-tuple/list or an Nxk array/tensor,"
" where k == 4 or 5")
arr = torch.tensor(box)[None, :]
else:
# avoid modifying the input box
if is_numpy:
arr = torch.from_numpy(np.asarray(box)).clone()
else:
arr = box.clone()
assert to_mode.value not in [
BoxMode.XYXY_REL,
BoxMode.XYWH_REL,
] and from_mode.value not in [
BoxMode.XYXY_REL,
BoxMode.XYWH_REL,
], "Relative mode not yet supported!"
if from_mode == BoxMode.XYWHA_ABS and to_mode == BoxMode.XYXY_ABS:
assert (
arr.shape[-1] == 5
), "The last dimension of input shape must be 5 for XYWHA format"
original_dtype = arr.dtype
arr = arr.double()
w = arr[:, 2]
h = arr[:, 3]
a = arr[:, 4]
c = torch.abs(torch.cos(a * math.pi / 180.0))
s = torch.abs(torch.sin(a * math.pi / 180.0))
# This basically computes the horizontal bounding rectangle of the rotated box
new_w = c * w + s * h
new_h = c * h + s * w
# convert center to top-left corner
arr[:, 0] -= new_w / 2.0
arr[:, 1] -= new_h / 2.0
# bottom-right corner
arr[:, 2] = arr[:, 0] + new_w
arr[:, 3] = arr[:, 1] + new_h
arr = arr[:, :4].to(dtype=original_dtype)
elif from_mode == BoxMode.XYWH_ABS and to_mode == BoxMode.XYWHA_ABS:
original_dtype = arr.dtype
arr = arr.double()
arr[:, 0] += arr[:, 2] / 2.0
arr[:, 1] += arr[:, 3] / 2.0
angles = torch.zeros((arr.shape[0], 1), dtype=arr.dtype)
arr = torch.cat((arr, angles), axis=1).to(dtype=original_dtype)
else:
if to_mode == BoxMode.XYXY_ABS and from_mode == BoxMode.XYWH_ABS:
arr[:, 2] += arr[:, 0]
arr[:, 3] += arr[:, 1]
elif from_mode == BoxMode.XYXY_ABS and to_mode == BoxMode.XYWH_ABS:
arr[:, 2] -= arr[:, 0]
arr[:, 3] -= arr[:, 1]
else:
raise NotImplementedError(
"Conversion from BoxMode {} to {} is not supported yet".
format(from_mode, to_mode))
if single_box:
return original_type(arr.flatten().tolist())
if is_numpy:
return arr.numpy()
else:
return arr
def cubic_spline(inputs,
unnormalized_widths,
unnormalized_heights,
unnorm_derivatives_left,
unnorm_derivatives_right,
inverse=False,
left=0.,
right=1.,
bottom=0.,
top=1.,
min_bin_width=DEFAULT_MIN_BIN_WIDTH,
min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
eps=DEFAULT_EPS,
quadratic_threshold=DEFAULT_QUADRATIC_THRESHOLD):
"""
References:
> Blinn, J. F. (2007). How to solve a cubic equation, part 5: Back to numerics. IEEE Computer
Graphics and Applications, 27(3):78–89.
"""
if not inverse and (torch.min(inputs) < left or torch.max(inputs) > right):
raise ValueError('Input is outside the domain')
elif inverse and (torch.min(inputs) < bottom or torch.max(inputs) > top):
raise ValueError('Input is outside the domain')
num_bins = unnormalized_widths.shape[-1]
if min_bin_width * num_bins > 1.0:
raise ValueError('Minimal bin width too large for the number of bins')
if min_bin_height * num_bins > 1.0:
raise ValueError('Minimal bin height too large for the number of bins')
if inverse:
inputs = (inputs - bottom) / (top - bottom)
else:
inputs = (inputs - left) / (right - left)
widths = F.softmax(unnormalized_widths, dim=-1)
widths = min_bin_width + (1 - min_bin_width * num_bins) * widths
cumwidths = torch.cumsum(widths, dim=-1)
cumwidths[..., -1] = 1
cumwidths = F.pad(cumwidths, pad=(1, 0), mode='constant', value=0.0)
heights = F.softmax(unnormalized_heights, dim=-1)
heights = min_bin_height + (1 - min_bin_height * num_bins) * heights
cumheights = torch.cumsum(heights, dim=-1)
cumheights[..., -1] = 1
cumheights = F.pad(cumheights, pad=(1, 0), mode='constant', value=0.0)
slopes = heights / widths
min_something_1 = torch.min(torch.abs(slopes[..., :-1]),
torch.abs(slopes[..., 1:]))
min_something_2 = (0.5 * (widths[..., 1:] * slopes[..., :-1] +
widths[..., :-1] * slopes[..., 1:]) /
(widths[..., :-1] + widths[..., 1:]))
min_something = torch.min(min_something_1, min_something_2)
derivatives_left = torch.sigmoid(unnorm_derivatives_left) * 3 * slopes[
..., 0][..., None]
derivatives_right = torch.sigmoid(unnorm_derivatives_right) * 3 * slopes[
..., -1][..., None]
derivatives = min_something * (torch.sign(slopes[..., :-1]) +
torch.sign(slopes[..., 1:]))
derivatives = torch.cat([derivatives_left, derivatives, derivatives_right],
dim=-1)
a = (derivatives[..., :-1] + derivatives[..., 1:] -
2 * slopes) / widths.pow(2)
b = (3 * slopes - 2 * derivatives[..., :-1] -
derivatives[..., 1:]) / widths
c = derivatives[..., :-1]
d = cumheights[..., :-1]
if inverse:
bin_idx = nf.util.searchsorted(cumheights, inputs)[..., None]
else:
bin_idx = nf.util.searchsorted(cumwidths, inputs)[..., None]
inputs_a = a.gather(-1, bin_idx)[..., 0]
inputs_b = b.gather(-1, bin_idx)[..., 0]
inputs_c = c.gather(-1, bin_idx)[..., 0]
inputs_d = d.gather(-1, bin_idx)[..., 0]
input_left_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0]
input_right_cumwidths = cumwidths.gather(-1, bin_idx + 1)[..., 0]
if inverse:
# Modified coefficients for solving the cubic.
inputs_b_ = (inputs_b / inputs_a) / 3.
inputs_c_ = (inputs_c / inputs_a) / 3.
inputs_d_ = (inputs_d - inputs) / inputs_a
delta_1 = -inputs_b_.pow(2) + inputs_c_
delta_2 = -inputs_c_ * inputs_b_ + inputs_d_
delta_3 = inputs_b_ * inputs_d_ - inputs_c_.pow(2)
discriminant = 4. * delta_1 * delta_3 - delta_2.pow(2)
depressed_1 = -2. * inputs_b_ * delta_1 + delta_2
depressed_2 = delta_1
three_roots_mask = discriminant > 0 # Discriminant == 0 might be a problem in practice.
one_root_mask = discriminant <= 0
outputs = torch.zeros_like(inputs)
# Deal with one root cases.
p = cbrt((-depressed_1[one_root_mask] +
torch.sqrt(-discriminant[one_root_mask])) / 2.)
q = cbrt((-depressed_1[one_root_mask] -
torch.sqrt(-discriminant[one_root_mask])) / 2.)
outputs[one_root_mask] = ((p + q) - inputs_b_[one_root_mask] +
input_left_cumwidths[one_root_mask])
# Deal with three root cases.
theta = torch.atan2(torch.sqrt(discriminant[three_roots_mask]),
-depressed_1[three_roots_mask])
theta /= 3.
cubic_root_1 = torch.cos(theta)
cubic_root_2 = torch.sin(theta)
root_1 = cubic_root_1
root_2 = -0.5 * cubic_root_1 - 0.5 * math.sqrt(3) * cubic_root_2
root_3 = -0.5 * cubic_root_1 + 0.5 * math.sqrt(3) * cubic_root_2
root_scale = 2 * torch.sqrt(-depressed_2[three_roots_mask])
root_shift = (-inputs_b_[three_roots_mask] +
input_left_cumwidths[three_roots_mask])
root_1 = root_1 * root_scale + root_shift
root_2 = root_2 * root_scale + root_shift
root_3 = root_3 * root_scale + root_shift
root1_mask = ((input_left_cumwidths[three_roots_mask] - eps) <
root_1).float()
root1_mask *= (
root_1 < (input_right_cumwidths[three_roots_mask] + eps)).float()
root2_mask = ((input_left_cumwidths[three_roots_mask] - eps) <
root_2).float()
root2_mask *= (
root_2 < (input_right_cumwidths[three_roots_mask] + eps)).float()
root3_mask = ((input_left_cumwidths[three_roots_mask] - eps) <
root_3).float()
root3_mask *= (
root_3 < (input_right_cumwidths[three_roots_mask] + eps)).float()
roots = torch.stack([root_1, root_2, root_3], dim=-1)
masks = torch.stack([root1_mask, root2_mask, root3_mask], dim=-1)
mask_index = torch.argsort(masks, dim=-1,
descending=True)[..., 0][..., None]
outputs[three_roots_mask] = torch.gather(roots,
dim=-1,
index=mask_index).view(-1)
# Deal with a -> 0 (almost quadratic) cases.
quadratic_mask = inputs_a.abs() < quadratic_threshold
a = inputs_b[quadratic_mask]
b = inputs_c[quadratic_mask]
c = (inputs_d[quadratic_mask] - inputs[quadratic_mask])
alpha = (-b + torch.sqrt(b.pow(2) - 4 * a * c)) / (2 * a)
outputs[quadratic_mask] = alpha + input_left_cumwidths[quadratic_mask]
shifted_outputs = (outputs - input_left_cumwidths)
logabsdet = -torch.log((3 * inputs_a * shifted_outputs.pow(2) +
2 * inputs_b * shifted_outputs + inputs_c))
else:
shifted_inputs = (inputs - input_left_cumwidths)
outputs = (inputs_a * shifted_inputs.pow(3) +
inputs_b * shifted_inputs.pow(2) +
inputs_c * shifted_inputs + inputs_d)
logabsdet = torch.log((3 * inputs_a * shifted_inputs.pow(2) +
2 * inputs_b * shifted_inputs + inputs_c))
if inverse:
outputs = outputs * (right - left) + left
logabsdet = logabsdet - math.log(top - bottom) + math.log(right - left)
else:
outputs = outputs * (top - bottom) + bottom
logabsdet = logabsdet + math.log(top - bottom) - math.log(right - left)
return outputs, logabsdet
def _inverse_stft(transform: Union[th.Tensor, Tuple[th.Tensor, th.Tensor]],
kernel: th.Tensor,
window: th.Tensor,
input: str = "polar",
frame_hop: int = 256,
onesided: bool = False,
center: bool = False) -> th.Tensor:
"""
iSTFT inner function
Args:
transform (Tensor or [Tensor, Tensor]), STFT transform results
kernel (Tensor), STFT transform kernels, from init_kernel(...)
input (str), input format:
polar: return (magnitude, phase) pair
complex: return (real, imag) pair
real: return [real; imag] Tensor
frame_hop: frame hop size in number samples
onesided: return half FFT bins
center: used in _forward_stft
Return:
wav (Tensor), N x S
"""
if input not in ["polar", "complex", "real"]:
raise ValueError(f"Unknown output format: {input}")
if input == "real":
real, imag = transform[..., 0], transform[..., 1]
elif input == "polar":
real = transform[0] * th.cos(transform[1])
imag = transform[0] * th.sin(transform[1])
else:
real, imag = transform
# (N) x F x T
imag_dim = imag.dim()
if imag_dim not in [2, 3]:
raise RuntimeError(f"Expect 2D/3D tensor, but got {imag_dim}D")
# if F x T, reshape 1 x F x T
if imag_dim == 2:
real = th.unsqueeze(real, 0)
imag = th.unsqueeze(imag, 0)
if onesided:
# [self.num_bins - 2, ..., 1]
reverse = range(kernel.shape[0] // 4 - 1, 0, -1)
# extend matrix: N x B x T
real = th.cat([real, real[:, reverse]], 1)
imag = th.cat([imag, -imag[:, reverse]], 1)
# pack: N x 2B x T
packed = th.cat([real, imag], dim=1)
# N x 1 x T
s = tf.conv_transpose1d(packed, kernel, stride=frame_hop, padding=0)
# normalized audio samples
# refer: https://github.com/pytorch/audio/blob/2ebbbf511fb1e6c47b59fd32ad7e66023fa0dff1/torchaudio/functional.py#L171
# 1 x W x T
win = th.repeat_interleave(window[None, ..., None],
packed.shape[-1],
dim=-1)
# W x 1 x W
I = th.eye(window.shape[0], device=win.device)[:, None]
# 1 x 1 x T
norm = tf.conv_transpose1d(win**2, I, stride=frame_hop, padding=0)
if center:
pad = kernel.shape[-1] // 2
s = s[..., pad:-pad]
norm = norm[..., pad:-pad]
s = s / (norm + EPSILON)
# N x S
s = s.squeeze(1)
return s
def log_prob(self, value):
log_prob = self.concentration * torch.cos(value - self.loc)
log_prob = log_prob - math.log(2 * math.pi) - _log_modified_bessel_fn(
self.concentration, order=0)
return log_prob
def reconstruct(self, depth, frame='w'):
"""
Reconstructs pixel-wise 3D points from a depth map.
Parameters
----------
depth : torch.Tensor [B,1,H,W]
Depth map for the camera
frame : 'w'
Reference frame: 'c' for camera and 'w' for world
Returns
-------
points : torch.tensor [B,3,H,W]
Pixel-wise 3D points
"""
B, C, H, W = depth.shape
assert C == 1
device = depth.get_device()
xi, yi = meshgrid(B, H, W, depth.dtype, depth.device, normalized=False)
xi = ((xi - (W - 1) / 2 - self.principal_point[:, 0].unsqueeze(1).unsqueeze(2).repeat([1, H, W]))
* self.scale_factors[:, 0].unsqueeze(1).unsqueeze(2).repeat([1, H, W])).unsqueeze(1)
yi = ((yi - (H - 1) / 2 - self.principal_point[:, 1].unsqueeze(1).unsqueeze(2).repeat([1, H, W]))
* self.scale_factors[:, 1].unsqueeze(1).unsqueeze(2).repeat([1, H, W])).unsqueeze(1)
use_precomputed_theta_lut = False
if use_precomputed_theta_lut:
theta_tensor = torch.zeros(B, 1, H, W)
for b in range(B):
theta_tensor[b, 0] = torch.from_numpy(np.load(self.path_to_theta_lut[b]))
theta_tensor = theta_tensor.to(device)
else:
N = 12
theta_tensor = (torch.zeros(B, 1, H, W)).to(device)
ri = torch.sqrt(xi.pow(2) + yi.pow(2))
for _ in range(N):
t1 = theta_tensor
t2 = theta_tensor * t1
t3 = theta_tensor * t2
t4 = theta_tensor * t3
theta_tensor = t1 + .5 * (ri - (self.poly_coeffs[:, 0].view(B, 1, 1, 1) * t1
+ self.poly_coeffs[:, 1].view(B, 1, 1, 1) * t2
+ self.poly_coeffs[:, 2].view(B, 1, 1, 1) * t3
+ self.poly_coeffs[:, 3].view(B, 1, 1, 1) * t4)) \
/ (self.poly_coeffs[:, 0].view(B, 1, 1, 1)
+ 2 * self.poly_coeffs[:, 1].view(B, 1, 1, 1) * t1
+ 3 * self.poly_coeffs[:, 2].view(B, 1, 1, 1) * t2
+ 4 * self.poly_coeffs[:, 3].view(B, 1, 1, 1) * t3)
# l'astuce pour que ça marche a été de multiplier la mise à jour par 0.5 (au lieu de 1 selon Newton...)
#get_roots_table_tensor(self.poly_coeffs, self.principal_point, self.scale_factors, H, W).to(device)
rc = depth * torch.sin(theta_tensor)
#yi, xi = centered_2d_grid(B, H, W, depth.dtype, depth.device, self.principal_point, self.scale_factors)
phi = torch.atan2(yi, xi).to(device)
xc = rc * torch.cos(phi)
yc = rc * torch.sin(phi)
zc = depth * torch.cos(theta_tensor)
#print(zc[0, 0, :, 127])
# mask = (depth == 0).detach()
# xc[mask] = 0.
# yc[mask] = 0.
# zc[mask] = 0.
Xc = torch.cat([xc, yc, zc], dim=1)
# If in camera frame of reference
if frame == 'c':
return Xc
# If in world frame of reference
elif frame == 'w':
return self.Twc @ Xc
# If none of the above
else:
raise ValueError('Unknown reference frame {}'.format(frame))
def warmup_cosine(x, warmup=0.002):
if x < warmup:
return x / warmup
return 0.5 * (1.0 + torch.cos(math.pi * x))
def forward(self, v):
x = self.base
# Project v onto {<x,v> = 0}
v = self.frame(x, v)
vnorm = v.norm(dim=-1, keepdim=True)
return self.radius * (torch.cos(vnorm) * x + sinc(vnorm) * v)
def forward(self, event_times):
N, T = event_times.shape
pe = torch.zeros(N, T, self.dim).to(event_times)
pe[:, :, 0::2] = torch.sin(event_times[..., None] * self.div_term)
pe[:, :, 1::2] = torch.cos(event_times[..., None] * self.div_term)
return pe
def reference(ref_data):
a = torch.tensor(d1, requires_grad=True)
b = torch.cos(a)
d__o = ref_data.getOutputTensorGrad(0)
b.backward(torch.tensor(d__o))
return [b, a.grad, None]
def phase_vocoder(complex_specgrams, rate, phase_advance):
# type: (Tensor, float, Tensor) -> Tensor
r"""Given a STFT tensor, speed up in time without modifying pitch by a
factor of ``rate``.
Args:
complex_specgrams (torch.Tensor): Dimension of `(channel, freq, time, complex=2)`
rate (float): Speed-up factor
phase_advance (torch.Tensor): Expected phase advance in each bin. Dimension
of (freq, 1)
Returns:
complex_specgrams_stretch (torch.Tensor): Dimension of `(channel,
freq, ceil(time/rate), complex=2)`
Example
>>> freq, hop_length = 1025, 512
>>> # (channel, freq, time, complex=2)
>>> complex_specgrams = torch.randn(2, freq, 300, 2)
>>> rate = 1.3 # Speed up by 30%
>>> phase_advance = torch.linspace(
>>> 0, math.pi * hop_length, freq)[..., None]
>>> x = phase_vocoder(complex_specgrams, rate, phase_advance)
>>> x.shape # with 231 == ceil(300 / 1.3)
torch.Size([2, 1025, 231, 2])
"""
time_steps = torch.arange(0,
complex_specgrams.size(-2),
rate,
device=complex_specgrams.device,
dtype=complex_specgrams.dtype)
alphas = time_steps % 1.0
phase_0 = angle(complex_specgrams[:, :, :1])
# Time Padding
complex_specgrams = torch.nn.functional.pad(complex_specgrams,
[0, 0, 0, 2])
# (new_bins, freq, 2)
complex_specgrams_0 = complex_specgrams[:, :, time_steps.long()]
complex_specgrams_1 = complex_specgrams[:, :, (time_steps + 1).long()]
angle_0 = angle(complex_specgrams_0)
angle_1 = angle(complex_specgrams_1)
norm_0 = torch.norm(complex_specgrams_0, p=2, dim=-1)
norm_1 = torch.norm(complex_specgrams_1, p=2, dim=-1)
phase = angle_1 - angle_0 - phase_advance
phase = phase - 2 * math.pi * torch.round(phase / (2 * math.pi))
# Compute Phase Accum
phase = phase + phase_advance
phase = torch.cat([phase_0, phase[:, :, :-1]], dim=-1)
phase_acc = torch.cumsum(phase, -1)
mag = alphas * norm_1 + (1 - alphas) * norm_0
real_stretch = mag * torch.cos(phase_acc)
imag_stretch = mag * torch.sin(phase_acc)
complex_specgrams_stretch = torch.stack([real_stretch, imag_stretch],
dim=-1)
return complex_specgrams_stretch
def reference(ref_data):
a = torch.tensor(d1, requires_grad=True)
b = torch.cos(a)
return [b]
import random
import torch
import unittest
import gpytorch
from gpytorch.kernels import RBFKernel, MultitaskKernel, GridInterpolationKernel
from gpytorch.means import ConstantMean, MultitaskMean
from gpytorch.likelihoods import MultitaskGaussianLikelihood
from gpytorch.random_variables import MultitaskGaussianRandomVariable
# Simple training data: let's try to learn a sine function
train_x = torch.linspace(0, 1, 100)
# y1 function is sin(2*pi*x) with noise N(0, 0.04)
train_y1 = torch.sin(train_x * (2 * pi)) + torch.randn(train_x.size()) * 0.1
# y2 function is cos(2*pi*x) with noise N(0, 0.04)
train_y2 = torch.cos(train_x * (2 * pi)) + torch.randn(train_x.size()) * 0.1
# Create a train_y which interleaves the two
train_y = torch.stack([train_y1, train_y2], -1)
class MultitaskGPModel(gpytorch.models.ExactGP):
def __init__(self, train_x, train_y, likelihood):
super(MultitaskGPModel, self).__init__(train_x, train_y, likelihood)
self.mean_module = MultitaskMean(ConstantMean(), n_tasks=2)
self_covar_module = GridInterpolationKernel(RBFKernel(),
grid_size=100,
grid_bounds=[(0, 1)])
self.covar_module = MultitaskKernel(self_covar_module,
n_tasks=2,
rank=1)
def sinusoidal_encode(x, w):
y = w * x
y[1:, 0::2] = torch.sin(y[1:, 0::2].clone())
y[1:, 1::2] = torch.cos(y[1:, 1::2].clone())
return y
'l1_normalization':
nnef_l1_normalization,
'l2_normalization':
nnef_l2_normalization,
'batch_normalization':
nnef_batch_normalization,
# 'avg_roi_pool': unsupported,
# 'max_roi_pool': unsupported,
# 'roi_resample': unsupported,
# 'avg_roi_align': unsupported,
# 'max_roi_align': unsupported,
'linear_quantize':
nnef_linear_quantize,
'logarithmic_quantize':
nnef_logarithmic_quantize,
'copy_n':
nnef_copy_n,
'sin':
lambda x: torch.sin(x),
'cos':
lambda x: torch.cos(x),
'tile':
lambda input, repeats: input.repeat(*repeats),
'pad':
nnef_pad,
'any_reduce':
lambda input, axes: _nnef_generic_reduce(input, axes=axes, f=torch.any),
'all_reduce':
lambda input, axes: _nnef_generic_reduce(input, axes=axes, f=torch.all),
}
def psf_func(self, X_os, Y_os, Z, I, add_wmap=True):
if add_wmap:
self.filt_size = self.w_map.shape[-1]
v = torch.arange(self.filt_size) - self.filt_size // 2
v = v.reshape([1, self.filt_size]).float().cuda()
if self.psf_pars['modality'] == 'GAUSS_2D':
W_x = (v[None, :, :] - X_os)**2
W_y = (v.transpose(1, 0)[None, :, :] - Y_os)**2
W = self.psf_pars['w_prop'] * torch.exp(
-(W_x + W_y) /
(2 * (Z * self.psf_pars['width1'])**2 + self.taylor_corr)) / (
2 * np.pi *
(Z * self.psf_pars['width1'])**2 + self.taylor_corr)
W += (1 - self.psf_pars['w_prop']) * torch.exp(
-(W_x + W_y) /
(2 * (Z * self.psf_pars['width2'])**2 + self.taylor_corr)) / (
2 * np.pi *
(Z * self.psf_pars['width2'])**2 + self.taylor_corr)
if self.psf_pars['modality'] == 'ASTIG_3D':
om_x = 0.5 * self.psf_pars['om_0'] * torch.sqrt(1 + (
(Z - self.psf_pars['c_x']) / self.psf_pars['d'])**2)
om_y = 0.5 * self.psf_pars['om_0'] * torch.sqrt(1 + (
(Z - self.psf_pars['c_y']) / self.psf_pars['d'])**2)
W_x = (v[None, :, :] - X_os)**2
W_y = (v.transpose(1, 0)[None, :, :] - Y_os)**2
W = torch.exp(-W_x /
(2 * (om_x**2 + self.taylor_corr))) * torch.exp(
-W_y / (2 * (om_y**2 + self.taylor_corr))) / (
2 * np.pi * (om_x * om_y + self.taylor_corr))
if self.psf_pars['modality'] == 'HELIX_3D':
x_shift = -self.psf_pars['rad'] * torch.cos(
self.psf_pars['rot_multi'] * Z + self.psf_pars['zero_rot'])
y_shift = self.psf_pars['rad'] * torch.sin(
self.psf_pars['rot_multi'] * Z + self.psf_pars['zero_rot'])
W1 = torch.sqrt((v[None, :, :] - (X_os + x_shift))**2 +
(v.transpose(1, 0)[None, :, :] -
(Y_os + y_shift))**2)
W = torch.exp(-W1**2 /
(2 *
(self.psf_pars['width'])**2 + self.taylor_corr)) / (
2 * np.pi *
(self.psf_pars['width'])**2 + self.taylor_corr)
W2 = torch.sqrt((v[None, :, :] - (X_os - x_shift))**2 +
(v.transpose(1, 0)[None, :, :] -
(Y_os - y_shift))**2)
W += torch.exp(
-W2**2 /
(2 * (self.psf_pars['width'])**2 + self.taylor_corr)) / (
2 * np.pi * (self.psf_pars['width'])**2 + self.taylor_corr)
W /= 2
if self.psf_pars['modality'] == 'ELIPSE_2D':
om_x = Z * self.psf_pars['om_0'] * (1 +
self.psf_pars['ellipticity'])
om_y = Z * self.psf_pars['om_0']
W_x = (v[None, :, :] - X_os)**2
W_y = (v.transpose(1, 0)[None, :, :] - Y_os)**2
W = torch.exp(-W_x /
(2 * (om_x**2 + self.taylor_corr))) * torch.exp(
-W_y / (2 * (om_y**2 + self.taylor_corr))) / (
2 * np.pi * (om_x * om_y + self.taylor_corr))
if add_wmap:
maps = self.interpolate_tri(self.w_map, X_os, Y_os,
self.z_for_map_ind, self.interpolate)
W += maps
W /= W.sum(-1).sum(-1)[:, None, None]
W *= I
return W
def forward(self, wave, gt, fps, pred = None, flag = None): # all variable operation
if flag is not None:
idx = flag.eq(1);
wave = wave[idx,:];
gt = gt[idx,:];
fps = fps[idx,:];
pred = pred[idx,:];
if(gt.shape[0] == 0):
loss = 0.0;
return loss, 0;
hr = torch.mul(gt, fps);
hr = hr*60/self.clip_length;
hr[hr.ge(self.high_bound)] = self.high_bound-1;
hr[hr.le(self.low_bound)] = self.low_bound;
if pred is not None:
pred = torch.mul(pred, fps);
pred = pred * 60 / self.clip_length;
batch_size = wave.shape[0];
f_t = self.bpm_range / fps;
preds = wave * self.hanning;
preds = preds.view(batch_size, 1, -1);
f_t = f_t.view(batch_size, -1, 1);
tmp = self.two_pi_n.repeat(batch_size, 1);
tmp = tmp.view(batch_size, 1, -1)
complex_absolute = torch.sum(preds * torch.sin(f_t*tmp), dim=-1) ** 2 \
+ torch.sum(preds * torch.cos(f_t*tmp), dim=-1) ** 2
target = hr - self.low_bound;
target = target.type(torch.long).view(batch_size);
whole_max_val, whole_max_idx = complex_absolute.max(1)
whole_max_idx = whole_max_idx + self.low_bound;
if self.loss_type == 1:
loss = self.cross_entropy(complex_absolute, target);
elif self.loss_type == 7:
norm_t = (torch.ones(batch_size).cuda() / torch.sum(complex_absolute, dim = 1));
norm_t = norm_t.view(-1,1);
complex_absolute = complex_absolute * norm_t;
loss = self.cross_entropy(complex_absolute, target);
idx_l = target - self.delta;
idx_l[idx_l.le(0)] = 0;
idx_r = target + self.delta;
idx_r[idx_r.ge(self.high_bound - self.low_bound - 1)] = self.high_bound - self.low_bound - 1;
loss_snr = 0.0;
for i in range(0, batch_size):
loss_snr = loss_snr + 1 - torch.sum(complex_absolute[i, idx_l[i]:idx_r[i]]);
loss_snr = loss_snr / batch_size;
loss = loss + loss_snr;
return loss, whole_max_idx
def makeProjection(self, gti_b, inj_b, trs_b, rot_b, sca_b):
Tensor = torch.cuda.FloatTensor
n_batches = gti_b.shape[0]
projection = Variable(
Tensor(np.zeros((n_batches, self.window_size, self.window_size))))
for batch in range(n_batches):
conn = gti_b[batch].clone()
conn = conn.cpu().numpy()
conn = np.uint16(measure.label(conn, background=0))
n_instances = np.amax(conn)
trs = trs_b[batch]
rot = rot_b[batch]
sca = sca_b[batch]
inj = inj_b[batch]
for ins in range(1, n_instances + 1):
indices = np.argwhere(conn == ins)
ins_mask = np.zeros((self.window_size, self.window_size))
ins_mask[indices[:, 0], indices[:, 1]] = 1.0
ins_mask = Variable(Tensor(ins_mask))
# Compute center of mass
bx = np.mean(indices[:, 1])
by = np.mean(indices[:, 0])
remove = torch.mean(inj[indices[:, 0], indices[:, 1]])
if remove < 0.5: # if close to zero the instance is not to remove
ti = torch.mean(trs[0, indices[:, 0], indices[:, 1]])
tj = torch.mean(trs[1, indices[:, 0], indices[:, 1]])
r = torch.mean(rot[0, indices[:, 0], indices[:, 1]])
s = torch.mean(sca[0, indices[:, 0], indices[:, 1]])
# Computation of the homograpy
bx = (
(self.window_size // 2) - bx) / (self.window_size // 2)
by = (
(self.window_size // 2) - by) / (self.window_size // 2)
R = torch.eye(3, 3)
R[0, 0] = torch.cos(r)
R[0, 1] = -torch.sin(r)
R[1, 0] = torch.sin(r)
R[1, 1] = torch.cos(r)
T = torch.eye(3, 3)
T[0, 2] = ti
T[1, 2] = tj
S = torch.eye(3, 3)
S[0, 0] = 1 + s
S[1, 1] = 1 + s
B = torch.eye(3, 3)
B[0, 2] = bx
B[1, 2] = by
B_ = torch.eye(3, 3)
B_[0, 2] = -bx
B_[1, 2] = -by
H = torch.mm(R, B)
H = torch.mm(S, H)
H = torch.mm(B_, H)
H = torch.mm(T, H)
H = H.inverse().cuda()
#H[0,0] = 1.0
#H[1,1] = 1.0
#H[2,2] = 1.0
#H[0,2] = ti
#H[1,2] = tj
ins_mask = self.warper(
ins_mask.view(1, 1, self.window_size,
self.window_size), H.view(1, 1, 3, 3))
ins_mask = ins_mask[0, 0, :, :]
projection[batch] += ins_mask
return projection
def register_fn(model):
for args in init_args:
test_example = Example(model, args)
TEST_EXAMPLES.append(test_example)
EXAMPLE_IDS.append(model.__name__)
return register_fn
@register_model([
ModelArgs(
step_size=0.01,
num_steps=100,
q_i={'x': torch.tensor([0.0])},
p_i={'x': torch.tensor([1.0])},
q_f={'x': torch.sin(torch.tensor([1.0]))},
p_f={'x': torch.cos(torch.tensor([1.0]))},
prec=1e-4
)
])
class HarmonicOscillator(object):
@staticmethod
def energy(q, p):
return 0.5 * p['x'] ** 2 + 0.5 * q['x'] ** 2
@staticmethod
def potential_fn(q):
return 0.5 * q['x'] ** 2
@register_model([
ModelArgs(
