def test_min_with_index(seed, ctx, func_name, inshape, axis, keepdims):
x = np.random.RandomState(seed).randn(*inshape).astype(np.float32)
x = nn.Variable.from_numpy_array(x)
with nn.context_scope(ctx), nn.auto_forward(True):
val, idx = F.min(x, axis, keepdims, with_index=True)
assert_allclose(val.d, np.amin(x.d, axis, keepdims=keepdims))
shape = [a for i, a in enumerate(x.d.shape) if i not in axis] + [-1]
assert np.all(idx.d == x.d.reshape(*shape).argmin(-1).reshape(idx.d.shape))
with nn.context_scope(ctx), nn.auto_forward(True):
idx = F.min(x, axis, keepdims, only_index=True)
shape = [a for i, a in enumerate(x.d.shape) if i not in axis] + [-1]
assert np.all(idx.d == x.d.reshape(*shape).argmin(-1).reshape(idx.d.shape))
def chamfer_hausdorff_oneside_dists(X0, X1):
b0 = X0.shape[0]
b1 = X1.shape[0]
sum_ = 0
max_ = nn.NdArray.from_numpy_array(np.array(-np.inf))
n = 0
for i in tqdm.tqdm(range(0, b0, sub_batch_size),
desc="cdist-outer-loop"):
x0 = nn.NdArray.from_numpy_array(X0[i:i + sub_batch_size])
norm_x0 = F.sum(x0**2.0, axis=1, keepdims=True)
min_ = nn.NdArray.from_numpy_array(np.ones(x0.shape[0]) * np.inf)
for j in tqdm.tqdm(range(0, b1, sub_batch_size),
desc="cdist-inner-loop"):
x1 = nn.NdArray.from_numpy_array(X1[j:j + sub_batch_size])
# block pwd
norm_x1 = F.transpose(F.sum(x1**2.0, axis=1, keepdims=True),
(1, 0))
x1_T = F.transpose(x1, (1, 0))
x01 = F.affine(x0, x1_T)
bpwd = (norm_x0 + norm_x1 - 2.0 * x01)**0.5
# block min
min_ = F.minimum2(min_, F.min(bpwd, axis=1))
# sum/max over cols
sum_ += F.sum(min_)
n += bpwd.shape[0]
max_ = F.maximum2(max_, F.max(min_))
ocd = sum_.data / n
ohd = max_.data
return ocd, ohd
def forward_impl(self, inputs, outputs):
x = inputs[0].data
m = inputs[1].data
M = inputs[2].data
y = outputs[0].data
y.copy_from(x)
if not self.training:
return
mb = F.min(x, keepdims=True)
Mb = F.max(x, keepdims=True)
F.minimum2(m, mb, outputs=[m])
F.maximum2(M, Mb, outputs=[M])
def forward_impl(self, inputs, outputs):
x = inputs[0].data
m = inputs[1].data
M = inputs[2].data
y = outputs[0].data
y.copy_from(x)
if not self.training:
return
mb = F.min(x, keepdims=True)
Mb = F.max(x, keepdims=True)
F.identity(self.decay * m + (1 - self.decay) * mb, outputs=[m])
F.identity(self.decay * M + (1 - self.decay) * Mb, outputs=[M])
def ray_march(self, camloc, raydir, t0, t1, N, n_chunks, t_argmin=False):
# Points computation
BR, _ = t0.shape
t0 = F.reshape(t0, (BR, 1, 1))
t1 = F.reshape(t1, (BR, 1, 1))
camloc = F.reshape(camloc, (BR, 1, 3))
raydir = F.reshape(raydir, (BR, 1, 3))
step = (t1 - t0) / (N - 1)
intervals = F.reshape(F.arange(0, N), (1, N, 1))
ts = t0 + step * intervals
points = camloc + ts * raydir
points = F.reshape(points, (BR * N, 3))
# SDF computation
sdf_points = []
batch = (BR * N) // n_chunks
for r in range(0, BR * N, batch):
sdf_points.append(self.sdf(points[r:r + batch, :]))
sdf_points = F.reshape(F.concatenate(*sdf_points, axis=0), (BR, N, 1)) if n_chunks != 1 else \
F.reshape(sdf_points[0], (BR, N, 1))
# t_argmin computation
if t_argmin:
idx_min = F.min(sdf_points, axis=1, keepdims=True, only_index=True)
t_argmin = F.reshape(F.gather(ts, idx_min, axis=1, batch_dims=1),
(BR, 1))
return t_argmin
# Intersection check
points = F.reshape(points, (BR, N, 3))
sdf_pos = F.greater_equal_scalar(sdf_points[:, :-1, :], 0)
sdf_neg = F.less_equal_scalar(sdf_points[:, 1:, :], 0)
mask_hit = sdf_pos * sdf_neg
decreasing_consts = F.reshape(F.arange(N, 1, -1), (1, N - 1, 1))
vals = mask_hit * decreasing_consts
idx_max = F.max(vals, axis=1, only_index=True)
points = points[:, :-1, :]
x_hit = F.gather(points, idx_max, axis=1, batch_dims=1)
x_hit = F.reshape(x_hit, (BR, 3))
mask_hit = F.greater_scalar(F.sum(mask_hit, axis=1), 0)
mask_hit = F.reshape(mask_hit, (BR, 1))
x_hit_rm0 = x_hit
step = F.reshape(step, (BR, 1))
raydir = F.reshape(raydir, (BR, 3))
x_hit_rm1 = x_hit_rm0 + step * raydir
return x_hit_rm0, x_hit_rm1, mask_hit
def __call__(self, x, return_encoding_indices=False):
x = F.transpose(x, (0, 2, 3, 1))
x_flat = x.reshape((-1, self.embedding_dim))
x_flat_squared = F.broadcast(F.sum(x_flat**2, axis=1, keepdims=True),
(x_flat.shape[0], self.num_embedding))
emb_wt_squared = F.transpose(
F.sum(self.embedding_weight**2, axis=1, keepdims=True), (1, 0))
distances = x_flat_squared + emb_wt_squared - 2 * \
F.affine(x_flat, F.transpose(self.embedding_weight, (1, 0)))
encoding_indices = F.min(distances,
only_index=True,
axis=1,
keepdims=True)
encoding_indices.need_grad = False
quantized = F.embed(
encoding_indices.reshape(encoding_indices.shape[:-1]),
self.embedding_weight).reshape(x.shape)
if return_encoding_indices:
return encoding_indices, F.transpose(quantized, (0, 3, 1, 2))
encodings = F.one_hot(encoding_indices, (self.num_embedding, ))
e_latent_loss = F.mean(
F.squared_error(quantized.get_unlinked_variable(need_grad=False),
x))
q_latent_loss = F.mean(
F.squared_error(quantized,
x.get_unlinked_variable(need_grad=False)))
loss = q_latent_loss + self.commitment_cost * e_latent_loss
quantized = x + (quantized - x).get_unlinked_variable(need_grad=False)
avg_probs = F.mean(encodings, axis=0)
perplexity = F.exp(-F.sum(avg_probs * F.log(avg_probs + 1.0e-10)))
return loss, F.transpose(quantized,
(0, 3, 1, 2)), perplexity, encodings
