def _create_label(self, anchor, bbox):
# label: 1 is positive, 0 is negative, -1 is dont care
label = -jt.ones((anchor.shape[0], ), dtype="int32")
argmax_ious, max_ious, gt_argmax_ious = self._calc_ious(anchor, bbox)
# assign negative labels first so that positive labels can clobber them
label[max_ious < self.neg_iou_thresh] = 0
# positive label: for each gt, anchor with highest iou
label[gt_argmax_ious] = 1
# positive label: above threshold IOU
label[max_ious >= self.pos_iou_thresh] = 1
# subsample positive labels if we have too many
n_pos = int(self.pos_ratio * self.n_sample)
pos_index = jt.where(label == 1)[0]
if len(pos_index) > n_pos:
tmp_index = np.arange(0, pos_index.shape[0])
np.random.shuffle(tmp_index)
disable_index = tmp_index[:pos_index.shape[0] - n_pos]
disable_index = pos_index[disable_index]
label[disable_index] = -1
# subsample negative labels if we have too many
n_neg = self.n_sample - jt.sum(label == 1).item()
neg_index = jt.where(label == 0)[0]
if len(neg_index) > n_neg:
tmp_index = np.arange(0, neg_index.shape[0])
np.random.shuffle(tmp_index)
disable_index = tmp_index[:neg_index.shape[0] - n_neg]
disable_index = neg_index[disable_index]
label[disable_index] = -1
return argmax_ious, label
def execute(self, trans_points, cp, voxel, gridSize, weight=1):
if len(trans_points.shape) == 4:
trans_points = trans_points.squeeze(dim=-1)
nb = pointClosestCellIndex(trans_points)
idx = jt.matmul(
nb, jt.transform.to_tensor(jt.array([(gridSize**2), gridSize, 1])))
mask = (1 - voxel.view((-1), (gridSize**3)).gather(1, idx))
idx = idx.unsqueeze(2)
idx = idx.repeat(1, 1, 3)
mask = mask.unsqueeze(2).repeat(1, 1, 3)
closest_points = cp.gather(1, idx)
self.constant = weight
distance = (trans_points - closest_points)
distance = (distance * mask)
# self.save_for_backward(distance)
self.saved_tensors = distance
return (jt.mean(jt.sum(jt.sum(jt.pow(distance, 2), 2), 1)) * weight)
def log_sum_exp(x):
"""Utility function for computing log_sum_exp while determining
This will be used to determine unaveraged confidence loss across
all examples in a batch.
Args:
x (Variable(tensor)): conf_preds from conf layers
"""
x_max = x.data.max()
return jt.log(jt.sum(jt.exp(x - x_max), 1)) + x_max
def planesymTransform(sample, plane):
abc = plane[:, 0:3].unsqueeze(1).repeat(1, sample.shape[1], 1)
d = plane[:, 3].unsqueeze(1).unsqueeze(1).repeat(1, sample.shape[1], 1)
fenzi = (jt.sum((sample * abc), 2, True) + d)
fenmu = (jt.norm(plane[:, 0:3], 2, 1, True).unsqueeze(1).repeat(
1, sample.shape[1], 1) + 1e-05)
x = (2 * jt.divide(fenzi, fenmu))
y = jt.multiply(x.repeat(1, 1, 3), (abc / fenmu))
return (sample - y)
def execute(self):
self.renderer.transform.set_eyes_from_angles(2.732, 30, 140)
image = self.renderer(self.vertices,
self.faces,
self.textures,
metallic_textures=self.metallic_textures,
roughness_textures=self.roughness_textures)
loss = jt.sum((image - self.image_ref).sqr())
return loss
def execute(self, xyz1, xyz2, points1, points2):
"""
Input:
xyz1: input points position data, [B, C, N]
xyz2: sampled input points position data, [B, C, S]
points1: input points data, [B, D, N]
points2: input points data, [B, D, S]
Return:
new_points: upsampled points data, [B, D', N]
"""
# xyz1 = xyz1.permute(0, 2, 1)
# xyz2 = xyz2.permute(0, 2, 1)
# points2 = points2.permute(0, 2, 1)
B, N, C = xyz1.shape
_, S, _ = xyz2.shape
if S == 1:
interpolated_points = points2.repeat(1, N, 1)
else:
dists = square_distance(xyz1, xyz2)
idx, dists = jt.argsort(dists, dim=-1)
dists, idx = dists[:, :, :3], idx[:, :, :3] # [B, N, 3]
dist_recip = 1.0 / (dists + 1e-8)
norm = jt.sum(dist_recip, dim=2, keepdims=True)
weight = dist_recip / norm
interpolated_points = jt.sum(index_points(points2, idx) *
weight.view(B, N, 3, 1),
dim=2)
if points1 is not None:
# points1 = points1.permute(0, 2, 1)
new_points = concat([points1, interpolated_points], dim=-1)
else:
new_points = interpolated_points
new_points = new_points.permute(0, 2, 1)
# l = len(self.mlp_convs)
for i, conv in self.mlp_convs.layers.items():
# conv = self.mlp_convs[i]
bn = self.mlp_bns[i]
new_points = self.relu(bn(conv(new_points)))
return new_points.permute(0, 2, 1)
def chamfer_loss(pc1, pc2, reduction='mean', dims='BNC', bidirectional=False):
''' return the chamfer loss from pc1 to pc2.
:param pc1: input point cloud
:type pc1: jittor array
:param pc2: input point cloud
:type pc2: jittor array
:param reduction: reduction method in batches, can be 'mean', 'sum', or None. Default: 'mean'.
:type reduction: str, optional
:param dims: a string that represents each dimension, can be
'[BNC]' ([batch, number of points, xyz]), or
'[BCN]' ([batch, xyz, number of points]). Default: 'BNC'.
:type dims: str, optional
Example:
>>> import jittor as jt
>>> from jittor.loss3d import chamfer_loss
>>> jt.flags.use_cuda = True
>>> pc1 = jt.rand([10, 100, 3], dtype=jt.float32)
>>> pc2 = jt.rand([10, 100, 3], dtype=jt.float32)
>>> cf = chamfer_loss(pc1, pc2, dims='BNC', bidirectional=True)
>>> print('chamfer loss =', cf.item())
'''
if bidirectional:
return chamfer_loss(pc1, pc2, reduction, dims) + chamfer_loss(
pc2, pc1, reduction, dims)
assert dims in ['BNC', 'BCN']
if dims == 'BCN':
pc1, pc2 = pc1.permute(0, 2, 1), pc2.permute(0, 2, 1)
batch_size_1, N, _ = pc1.shape
batch_size_2, M, _ = pc2.shape
assert batch_size_1 == batch_size_2
batch_size = batch_size_1
idx = jt.code([batch_size, N],
'int32', [pc1, pc2],
cpu_src=cpu_src,
cuda_src=cuda_src)
nearest_pts = pc2.reindex([batch_size, idx.shape[1], 3],
['i0', '@e0(i0, i1)', 'i2'],
extras=[idx])
chamfer_distance = (((pc1 - nearest_pts)**2).sum(dim=-1)).sqrt()
if reduction is None:
return chamfer_distance
elif reduction == 'sum':
return jt.sum(chamfer_distance)
elif reduction == 'mean':
return jt.mean(chamfer_distance)
def integrator(raw, z_vals, rays_d, raw_noise_std=0, white_bkgd=False):
"""Transforms model's predictions to semantically meaningful values.
Args:
raw: [num_rays, num_samples along ray, 4]. Prediction from model.
z_vals: [num_rays, num_samples along ray]. Integration time.
rays_d: [num_rays, 3]. Direction of each ray.
Returns:
rgb_map: [num_rays, 3]. Estimated RGB color of a ray.
disp_map: [num_rays]. Disparity map. Inverse of depth map.
acc_map: [num_rays]. Sum of weights along each ray.
weights: [num_rays, num_samples]. Weights assigned to each sampled color.
depth_map: [num_rays]. Estimated distance to object.
"""
raw2alpha = lambda raw, dists, act_fn=jt.nn.relu: 1. - jt.exp(-act_fn(raw)
* dists)
dists = z_vals[..., 1:] - z_vals[..., :-1]
dists = jt.concat([
dists,
jt.array(np.array([1e10]).astype(np.float32)).expand(
dists[..., :1].shape)
], -1) # [N_rays, N_samples]
dists = dists * jt.norm(rays_d.unsqueeze(-2), p=2, dim=-1)
rgb = jt.sigmoid(raw[..., :3]) # [N_rays, N_samples, 3]
noise = 0.
if raw_noise_std > 0.:
noise = jt.init.gauss(raw[..., 3].shape, raw.dtype) * raw_noise_std
alpha = raw2alpha(raw[..., 3] + noise, dists) # [N_rays, N_samples]
weights = alpha * jt.cumprod(
jt.concat([jt.ones(
(alpha.shape[0], 1)), 1. - alpha + 1e-10], -1), -1)[:, :-1]
rgb_map = jt.sum(weights.unsqueeze(-1) * rgb, -2) # [N_rays, 3]
depth_map = jt.sum(weights * z_vals, -1)
disp_map = 1. / jt.maximum(1e-10 * jt.ones_like(depth_map),
depth_map / jt.sum(weights, -1))
acc_map = jt.sum(weights, -1)
if white_bkgd:
rgb_map = rgb_map + (1. - acc_map.unsqueeze(-1))
return rgb_map, disp_map, acc_map, weights, depth_map
def get_rays(H, W, focal, c2w, intrinsic=None):
i, j = jt.meshgrid(jt.linspace(0, W - 1, W), jt.linspace(0, H - 1, H))
i = i.t()
j = j.t()
if intrinsic is None:
dirs = jt.stack([(i - W * .5) / focal, (j - H * .5) / focal,
jt.ones_like(i)], -1).unsqueeze(-2)
else:
i += 0.5
j += 0.5
dirs = jt.stack([i, j, jt.ones_like(i)], -1).unsqueeze(-2)
dirs = jt.sum(dirs * intrinsic[:3, :3], -1).unsqueeze(-2)
# Rotate ray directions from camera frame to the world frame
rays_d = jt.sum(
dirs * c2w[:3, :3],
-1) # dot product, equals to: [c2w.dot(dir) for dir in dirs]
# Translate camera frame's origin to the world frame. It is the origin of all rays.
rays_o = c2w[:3, -1].expand(rays_d.shape)
return rays_o, rays_d
def GGX(N, H, roughness):
a = roughness * roughness
a2 = a * a
NdotH = nn.relu(jt.sum(N * H, dim=2))
NdotH2 = (NdotH * NdotH).unsqueeze(2)
num = a2
denom = (NdotH2 * (a2 - 1.0) + 1.0)
denom = 3.1415 * denom * denom
return num / denom
def square_distance(src, dst):
"""
Calculate Euclid distance between each two points.
src^T * dst = xn * xm + yn * ym + zn * zm；
sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
= sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst
Input:
src: source points, [B, N, C]
dst: target points, [B, M, C]
Output:
dist: per-point square distance, [B, N, M]
"""
B, N, _ = src.shape
_, M, _ = dst.shape
dist = -2 * jt.matmul(src, dst.permute(0, 2, 1))
dist += jt.sum(src**2, -1).view(B, N, 1)
dist += jt.sum(dst**2, -1).view(B, 1, M)
return dist
def calc_gradient_penalty(netD, real_data, generated_data):
LAMBDA = 10
b_size = real_data.shape[0]
alpha = jt.random([b_size, 1, 1, 1])
alpha = alpha.broadcast(real_data)
interpolated = ((alpha * real_data.data) +
((1 - alpha) * generated_data.data))
prob_interpolated = netD(interpolated)
gradients = jt.grad(prob_interpolated, interpolated)
gradients = jt.reshape(gradients, [b_size, -1])
gradients_norm = jt.sqrt((jt.sum((gradients**2), dim=1) + 1e-12))
return (LAMBDA * ((gradients_norm - 1)**2).mean())
def square_distance(tensor1, tensor2):
"""
Calculate Euclid distance between each two points.
tensor1^T * tensor2 = xn * xm + yn * ym + zn * zm；
sum(tensor1^2, dim=-1) = xn*xn + yn*yn + zn*zn;
sum(tensor2^2, dim=-1) = xm*xm + ym*ym + zm*zm;
dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
= sum(tensor1**2,dim=-1)+sum(tensor2**2,dim=-1)-2*tensor1^T*dst
Input:
tensor1: source points, [B, N, C]
tensor2: target points, [B, M, C]
Output:
dist: per-point square distance, [B, N, M]
"""
# print (src.size(), dst.size())
B, N, _ = tensor1.shape
_, M, _ = tensor2.shape
dist = -2 * jt.matmul(tensor1, tensor2.permute(0, 2, 1))
dist += jt.sum(tensor1**2, -1).view(B, N, 1)
dist += jt.sum(tensor2**2, -1).view(B, 1, M)
return dist
def execute(self, mesh, eyes=None):
if self.Gbuffer == "albedo":
return mesh
if self.Gbuffer == "normal" or self.Gbuffer == "depth":
mesh.textures = jt.ones_like(mesh.textures)
if self.light_mode == 'surface':
diffuseLight = jt.zeros(mesh.faces.shape)
specularLight = jt.zeros(mesh.faces.shape)
diffuseLight = self.ambient(diffuseLight)
for directional in self.directionals:
[diffuseLight, specularLight] = directional(
diffuseLight, specularLight, mesh.surface_normals,
(jt.sum(mesh.face_vertices, dim=2) / 3.0), eyes,
mesh.with_specular, mesh.metallic_textures,
mesh.roughness_textures)
if len(mesh.textures.shape) == 4:
mesh.textures = jt.clamp(
mesh.textures * diffuseLight.unsqueeze(2) +
jt.ones_like(mesh.textures) * specularLight.unsqueeze(2),
0.0, 1.0)
elif len(mesh.textures.shape) == 6:
mesh.textures = jt.clamp(
mesh.textures *
diffuseLight.unsqueeze(2).unsqueeze(2).unsqueeze(2) +
jt.ones_like(mesh.textures) *
specularLight.unsqueeze(2).unsqueeze(2).unsqueeze(2), 0.0,
1.0)
elif self.light_mode == 'vertex':
diffuseLight = jt.zeros(mesh.vertices.shape)
specularLight = jt.zeros(mesh.vertices.shape)
diffuseLight = self.ambient(diffuseLight)
for directional in self.directionals:
[diffuseLight, specularLight
] = directional(diffuseLight, specularLight,
mesh.vertex_normals, mesh.vertices, eyes,
mesh.with_specular, mesh.metallic_textures,
mesh.roughness_textures)
if len(mesh.textures.shape) == 4:
mesh.textures = jt.clamp(
mesh.textures * diffuseLight.unsqueeze(2) +
jt.ones_like(mesh.textures) * specularLight.unsqueeze(2),
0.0, 1.0)
elif len(mesh.textures.shape) == 6:
mesh.textures = jt.clamp(
mesh.textures *
diffuseLight.unsqueeze(2).unsqueeze(2).unsqueeze(2) +
jt.ones_like(mesh.textures) *
specularLight.unsqueeze(2).unsqueeze(2).unsqueeze(2), 0.0,
1.0)
return mesh
def lighting(faces,
textures,
intensity_ambient=0.5,
intensity_directional=0.5,
color_ambient=(1, 1, 1),
color_directional=(1, 1, 1),
direction=(0, 1, 0)):
bs, nf = faces.shape[:2]
# arguments
# make sure to convert all inputs to float tensors
color_ambient = jt.array(color_ambient, "float32")
color_directional = jt.array(color_directional, "float32")
direction = jt.array(direction, "float32")
if len(color_ambient.shape) == 1:
color_ambient = color_ambient.unsqueeze(0)
if len(color_directional.shape) == 1:
color_directional = color_directional.unsqueeze(0)
if len(direction.shape) == 1:
direction = direction.unsqueeze(0)
# create light
light = jt.zeros((bs, nf, 3), "float32")
# ambient light
if intensity_ambient != 0:
light += intensity_ambient * color_ambient.unsqueeze(1)
# directional light
if intensity_directional != 0:
faces = faces.reshape((bs * nf, 3, 3))
v10 = faces[:, 0] - faces[:, 1]
v12 = faces[:, 2] - faces[:, 1]
normals = jt.normalize(jt.cross(v10, v12), eps=1e-5)
normals = normals.reshape((bs, nf, 3))
if len(direction.shape) == 2:
direction = direction.unsqueeze(1)
cos = nn.relu(jt.sum(normals * direction, dim=2))
# may have to verify that the next line is correct
light += intensity_directional * (color_directional.unsqueeze(1) *
cos.unsqueeze(2))
# apply
light = light.unsqueeze(-2).unsqueeze(-2).unsqueeze(-2)
textures *= light
return textures
def display_lincomb(proto_data, masks):
out_masks = jt.matmul(proto_data, masks.t())
# out_masks = cfg.mask_proto_mask_activation(out_masks)
for kdx in range(1):
jdx = kdx + 0
import matplotlib.pyplot as plt
coeffs = masks[jdx].numpy()
idx = np.argsort(-np.abs(coeffs))
# plt.bar(list(range(idx.shape[0])), coeffs[idx])
# plt.show()
coeffs_sort = coeffs[idx]
arr_h, arr_w = (4, 8)
proto_h, proto_w, _ = proto_data.shape
arr_img = np.zeros([proto_h * arr_h, proto_w * arr_w])
arr_run = np.zeros([proto_h * arr_h, proto_w * arr_w])
test = jt.sum(proto_data, -1).numpy()
for y in range(arr_h):
for x in range(arr_w):
i = arr_w * y + x
if i == 0:
running_total = proto_data[:, :,
idx[i]].numpy() * coeffs_sort[i]
else:
running_total += proto_data[:, :, idx[i]].numpy(
) * coeffs_sort[i]
running_total_nonlin = running_total
if cfg.mask_proto_mask_activation == activation_func.sigmoid:
running_total_nonlin = (
1 / (1 + np.exp(-running_total_nonlin)))
arr_img[y * proto_h:(y + 1) * proto_h, x * proto_w:(x + 1) *
proto_w] = (proto_data[:, :, idx[i]] / jt.max(
proto_data[:, :, idx[i]])).numpy() * coeffs_sort[i]
arr_run[y * proto_h:(y + 1) * proto_h, x * proto_w:(x + 1) *
proto_w] = (running_total_nonlin > 0.5).astype(
np.float)
plt.imshow(arr_img)
plt.show()
# plt.imshow(arr_run)
# plt.show()
# plt.imshow(test)
# plt.show()
plt.imshow(out_masks[:, :, jdx].numpy())
plt.show()
def test_longest_dis_fuse(self):
x = jt.array(np.random.rand(1, 3, 224, 224).astype(np.float32))
loss = jt.sum(resnet_fake(x))
ps = jt.find_vars('resnet_fake')
gs = jt.grad(loss, ps)
jt.sync(gs)
# assert not alloc big tensor
g = jt.dump_all_graphs()
for s in g.nodes_info:
if not s.startswith("Var"):
continue
shape = s.split("[")[1].split("]")[0].split(",")
ptr = s.split("(")[1].split(")")[0].split(",")[-1]
if ptr != '0':
assert len(shape) <= 5, s
def compute_mask_prob(self, proposal_embed, proposal_margin, pixel_embed):
m_h, m_w = pixel_embed.shape[-2:]
obj_num = proposal_embed.shape[0]
pixel_embed = pixel_embed.transpose(1, 2, 0).unsqueeze(0).expand(
obj_num, -1, -1, -1)
proposal_embed = proposal_embed.view(obj_num, 1, 1,
-1).expand(-1, m_h, m_w, -1)
if self.fix_margin:
proposal_margin = proposal_margin.new_ones(obj_num, m_h,
m_w) * self.init_margin
else:
proposal_margin = proposal_margin.view(obj_num, 1,
1).expand(-1, m_h, m_w)
mask_var = jt.sum((pixel_embed - proposal_embed).sqr(), dim=3)
mask_prob = jt.exp(-mask_var * proposal_margin)
return mask_prob
def get_patches(fake_B, face_mask):
# [1,1,512,512]->[bs,1,11,11]
patches = []
weights = []
W2 = int(W/2)
t = np.random.randint(res,size=2)
for i in range(aa):
for j in range(aa):
p = fake_B[:,:,t[0]+i*W:t[0]+(i+1)*W,t[1]+j*W:t[1]+(j+1)*W]
whitenum = jt.sum(p>=0.0)
if whitenum < 1 or whitenum > W*W-1:
continue
patches.append(p)
weights.append(face_mask[:,:,t[0]+i*W+W2,t[1]+j*W+W2])
fake_B_patches = jt.contrib.concat(patches, dim=0)
conti_weights = jt.contrib.concat(weights, dim=0)+1 #0->1,1->2
return fake_B_patches, conti_weights
def farthest_point_sample(xyz, npoint):
"""
Input:
xyz: pointcloud data, [B, N, C]
npoint: number of samples
Return:
centroids: sampled pointcloud index, [B, npoint]
"""
#import ipdb; ipdb.set_trace()
#device = xyz.device
B, N, C = xyz.shape
centroids = jt.zeros((B, npoint))
distance = jt.ones((B, N)) * 1e10
farthest = np.random.randint(0, N, B, dtype='l')
batch_indices = np.arange(B, dtype='l')
farthest = jt.array(farthest)
batch_indices = jt.array(batch_indices)
# jt.sync_all(True)
# print (xyz.shape, farthest.shape, batch_indices.shape, centroids.shape, distance.shape)
for i in range(npoint):
centroids[:, i] = farthest
centroid = xyz[batch_indices, farthest, :]
centroid = centroid.view(B, 1, 3)
dist = jt.sum((xyz - centroid.repeat(1, N, 1))**2, 2)
mask = dist < distance
# distance = mask.ternary(distance, dist)
# print (mask.size())
if mask.sum().data[0] > 0:
distance[mask] = dist[mask] # bug if mask.sum() == 0
farthest = jt.argmax(distance, 1)[0]
# print (farthest)
# print (farthest.shape)
# B, N, C = xyz.size()
# sample_list = random.sample(range(0, N), npoint)
# centroids = jt.zeros((1, npoint))
# centroids[0,:] = jt.array(sample_list)
# centroids = centroids.view(1, -1).repeat(B, 1)
# x_center = x[:,sample_list, :]
return centroids
def computeMaskIntersection(gt, gtMask, pred, predMask):
"""
Compute intersection between GT instance and prediction.
Increase efficiency by computing elementwise product between masks
only inside the tight bounding box of the union of the prediction and
target masks.
"""
if gtMask is None or predMask is None:
return 0
assert gtMask.shape == predMask.shape
assert len(gtMask.shape) == len(predMask.shape) == 2
xmin, ymin, xmax, ymax = getUnionBox(gt["box"], pred["box"])
gtMask_crop = gtMask[ymin:ymax, xmin:xmax]
predMask_crop = predMask[ymin:ymax, xmin:xmax]
# elementwise AND
intersection = jt.sum(gtMask_crop * predMask_crop).numpy().item()
return intersection
def R1Penalty(self, real_img, height, alpha):
# TODO: use_loss_scaling, for fp16
# apply_loss_scaling = lambda x: x * torch.exp(x * torch.Tensor([np.float32(np.log(2.0))]).to(real_img.device))
apply_loss_scaling = lambda x: x * jt.exp(x * jt.array(
[np.float32(np.log(2.0))]))
# undo_loss_scaling = lambda x: x * torch.exp(-x * torch.Tensor([np.float32(np.log(2.0))]).to(real_img.device))
undo_loss_scaling = lambda x: x * jt.exp(-x * jt.array(
[np.float32(np.log(2.0))]))
# real_img = torch.autograd.Variable(real_img, requires_grad=True)
real_img = init.constant(real_img.shape, 'float32', real_img)
assert not real_img.is_stop_grad()
real_logit = self.dis(real_img, height, alpha)
# real_logit = apply_loss_scaling(torch.sum(real_logit))
# real_grads = torch.autograd.grad(outputs=real_logit, inputs=real_img,
# grad_outputs=torch.ones(real_logit.size()).to(real_img.device),
# create_graph=True, retain_graph=True)[0].view(real_img.size(0), -1)
real_grads = jt.grad(real_logit, real_img).view(real_img.size(0), -1)
# real_grads = undo_loss_scaling(real_grads)
# r1_penalty = torch.sum(torch.mul(real_grads, real_grads))
r1_penalty = jt.sum(jt.multiply(real_grads, real_grads))
return r1_penalty
def test(net, test_data, state):
net.eval()
loss_avg = 0.0
correct = 0
start_time = time.time()
for batch_idx, (data, target) in enumerate(test_data):
data, target = jt.array(data), jt.array(target)
# forward
output = net(data)
loss = jt.nn.cross_entropy_loss(output, target)
# accuracy
pred = jt.argmax(output, dim=1)[0]
correct += float(jt.sum(pred == target).data[0])
# test loss average
loss_avg += float(loss.data[0])
end_time = time.time()
fps = (len(test_data) * test_data.batch_size) / (end_time - start_time)
state['test_loss'] = loss_avg / len(test_data)
state['test_accuracy'] = correct / (len(test_data) * test_data.batch_size)
state['test_fps'] = fps
def execute(self, input, target):
bs_idx = jt.array(range(input.shape[0]))
ret = (-jt.log(nn.softmax(input, dim=1)))[bs_idx, target]
if self.reduction != None:
ret = jt.mean(ret) if self.reduction == 'mean' else jt.sum(ret)
return ret
def test_node_performance(self):
mode = os.environ.get("test_node_performance")
if mode==None or mode not in "12":
return
if mode=="1":
bc = lambda x: jt.broadcast(x, [1,1,1,1],[0,1,2])
rd = lambda x: jt.sum(x)
else:
bc = lambda x: jt.reindex(x, [1,1,1,1],["i0+i1+i2+i3"])
rd = lambda x: jt.reindex_reduce(x, "add", [1], ["i0+i1+i2+i3"])
if jt.compiler.is_debug: return
def run():
start_time = time.time()
fop_num = 10000
fop_input_num = (2, 3) # (i,j) -> [i,i+j] -> [2, 5]
# fop_output_num = (1, 0) # [1,1]
inner_op_num = (0, 3)
fop_type_num = 63 # how many different fuse op
input_queue_num = 15
queue = [1.0]*(input_queue_num+1)
x = get_xorshf96()
rand = lambda x, l, r: l+((x())&r)
ops = ["add", "subtract", "multiply", "divide"]
get_op = lambda x: ops[(x())&3]
for i in range(fop_num):
prev = bc(queue[rand(x,0,input_queue_num)])
y = get_xorshf96(x()&fop_type_num)
inum = rand(y, *fop_input_num)
q = [prev]
for i in range(inum-1):
n = bc(queue[rand(x,0,input_queue_num)])
prev = jt.binary(prev, n, get_op(y))
q.append(prev)
innum = rand(y,*inner_op_num)
for _ in range(innum):
j = rand(y,0,len(q)-1)
n = q[j]
prev = jt.binary(prev, n, get_op(y))
q[j] = prev
prev = rd(prev)
queue[rand(x,0,input_queue_num)] = prev
a = jt.array(0.0)
for x in queue:
a += x
LOG.i("build graph", time.time()-start_time, jt.liveness_info().values())
start_time = time.time()
a.sync()
LOG.i("execute", time.time()-start_time)
# debug mode: build(0.68), execute(0.44)
# normal mode: build(0.56), execute(0.25)
# cast opt: build(0.50), execute(0.25)
# dtype opt: build(0.49), execute(0.25)
# pyjt opt: build(0.48), execute(0.25)
# ns opt: build(0.46), execute(0.24)
# nv opt: build(0.42), execute(0.23)
# nv opt: build(0.415),execute(0.225)
# jit_key opt: build(0.415),execute(0.15)
# jit_key opt: build(0.415),execute(0.11)
# sv opt: build(0.42), execute(0.12)
# noded opt: build(0.42), execute(0.10)
# tcm opt: build(0.40), execute(0.10)
# mode2: reindex
# jit_key opt: build(0.46),execute(0.12)
# noded opt: build(0.44),execute(0.11)
# for i in range(20):
# run()
for i in range(20):
run()
import gc
gc.collect()
run()
def directional_lighting(diffuseLight,
specularLight,
normals,
light_intensity=0.5,
light_color=(1, 1, 1),
light_direction=(0, 1, 0),
positions=None,
eye=None,
with_specular=False,
metallic_textures=None,
roughness_textures=None,
Gbuffer="None",
transform=None):
eye = jt.array(eye, "float32")
light_color = jt.array(light_color, "float32")
light_direction = jt.normalize(jt.array(light_direction, "float32"), dim=0)
if len(light_color.shape) == 1:
light_color = light_color.unsqueeze(0)
if len(light_direction.shape) == 1:
light_direction = light_direction.unsqueeze(0)
cosine = nn.relu(jt.sum(normals * light_direction, dim=2))
if with_specular:
if len(metallic_textures.shape) == 4:
total = metallic_textures.shape[2] * 1.0
metallic_textures = jt.sum(metallic_textures, dim=2) / total
roughness_textures = jt.sum(roughness_textures, dim=2) / total
elif len(metallic_textures.shape) == 6:
total = metallic_textures.shape[2] * metallic_textures.shape[
3] * metallic_textures.shape[4] * 1.0
metallic_textures = jt.sum(metallic_textures, dim=2)
metallic_textures = jt.sum(metallic_textures, dim=2)
metallic_textures = jt.sum(metallic_textures, dim=2)
metallic_textures = metallic_textures / total
roughness_textures = jt.sum(roughness_textures, dim=2)
roughness_textures = jt.sum(roughness_textures, dim=2)
roughness_textures = jt.sum(roughness_textures, dim=2)
roughness_textures = roughness_textures / total
#Microfacet model
if with_specular and (eye is not None) and (positions is not None) and (
metallic_textures is not None) and (roughness_textures
is not None):
N = normals
if len(eye.shape) == 2:
eye = eye.unsqueeze(1)
V = jt.normalize(eye - positions, dim=2)
L = light_direction
H = jt.normalize(V + L, dim=2)
#Default Setting
metallic = metallic_textures
roughness = roughness_textures
F0 = jt.array((0.04, 0.04, 0.04), "float32")
albedo = jt.array((1.0, 1.0, 1.0), "float32")
F0 = F0.unsqueeze(0).unsqueeze(1) * (
1 - metallic) + albedo.unsqueeze(0).unsqueeze(1) * metallic
radiance = light_intensity * (light_color.unsqueeze(1) *
cosine.unsqueeze(2))
#Cook-Torrance BRDF
NDF = GGX(N, H, roughness)
G = GeometrySmith(N, V, L, roughness)
F = fresnelSchlick(nn.relu(jt.sum(H * V, dim=2)), F0)
KS = F
KD = 1.0 - KS
KD *= (1.0 - metallic)
diffuseLight += KD * radiance
numerator = NDF * G * F
denominator = (4.0 * nn.relu(jt.sum(N * V, dim=2)) *
nn.relu(jt.sum(N * L, dim=2))).unsqueeze(2)
specular = numerator / jt.clamp(denominator, 0.01)
specularLight += specular * radiance
else:
diffuseLight += light_intensity * (light_color.unsqueeze(1) *
cosine.unsqueeze(2))
if Gbuffer == "normal":
specularLight *= 0.0
diffuseLight = normals * 0.5 + 0.5
elif Gbuffer == "depth":
specularLight *= 0.0
viewpos = transform.tranpos(positions)
diffuseLight = viewpos / jt.max(viewpos[..., 2])
diffuseLight[..., 0] = viewpos[..., 2] / jt.max(viewpos[..., 2])
diffuseLight[..., 1] = viewpos[..., 2] / jt.max(viewpos[..., 2])
return [diffuseLight, specularLight]
def execute(self, input, target):
ret = jt.abs(input - target)
if self.reduction != None:
ret = jt.mean(ret) if self.reduction == 'mean' else jt.sum(ret)
return ret
def execute(self):
self.renderer.eye = nr.get_points_from_angles(2.732, 0, 90)
image = self.renderer(self.vertices, self.faces, mode='silhouettes')
loss = jt.sum((image - self.image_ref.unsqueeze(0)).sqr())
return loss
def execute(self, xyz1, xyz2, points1, points2):
"""
Input:
xyz1: input points position data, [B, C, N]
xyz2: sampled input points position data, [B, C, S]
points1: input points data, [B, D, N]
points2: input points data, [B, D, S]
Return:
new_points: upsampled points data, [B, D', N]
"""
# print ('xyz1.shape, xyz2.shape')
# print (xyz1.shape, xyz2.shape, points1.shape, points2.shape)
xyz1 = xyz1.permute(0, 2, 1)
xyz2 = xyz2.permute(0, 2, 1)
points1 = points1.permute(0, 2, 1)
points2 = points2.permute(0, 2, 1)
B, N, C = xyz1.shape
_, S, _ = xyz2.shape
# points2 = points2.permute(0, 2, 1)
# print (xyz1.shape, xyz2.shape)
dists = square_distance(xyz1, xyz2)
idx, dists = jt.argsort(dists, dim=-1)
dists, idx = dists[:, :, :3], idx[:, :, :3] # [B, N, 3]
dist_recip = 1.0 / (dists + 1e-8)
norm = jt.sum(dist_recip, dim=2, keepdims=True)
weight = dist_recip / norm
interpolated_points = jt.sum(index_points(points2, idx) *
weight.view(B, N, 3, 1),
dim=2)
# print ('interpolated_points shape', interpolated_points.shape)
xyz_density = compute_density(xyz1, self.bandwidth)
density_scale = self.densitynet(xyz_density)
new_xyz, new_points, grouped_xyz_norm, _, grouped_density = sample_and_group(
N, self.nsample, xyz1, interpolated_points,
density_scale.reshape(B, N, 1))
new_points = new_points.permute(0, 3, 2, 1) # [B, C+D, nsample,npoint]
for i in range(len(self.mlp_convs)):
conv = self.mlp_convs[i]
bn = self.mlp_bns[i]
# print ('new new new point shape', new_points.shape)
new_points = self.relu(bn(conv(new_points)))
grouped_xyz = grouped_xyz_norm.permute(0, 3, 2, 1)
weights = self.weightnet(grouped_xyz)
new_points = new_points * grouped_density.permute(0, 3, 2, 1)
new_points = jt.matmul(new_points.permute(0, 3, 1, 2),
weights.permute(0, 3, 2, 1)).reshape(B, N, -1)
new_points = self.linear(new_points)
new_points = self.bn_linear(new_points.permute(0, 2, 1))
new_points = self.relu(new_points)
new_xyz = new_xyz.permute(0, 2, 1)
return new_points
def entropy(self):
return -jt.sum(jt.mean(self.probs) * jt.log(self.probs))