def seq_rnn_embed(vxs, exs, birnn, return_seqs=False): """Embed given sequences using rnn.""" # vxs.shape == (..., S) # exs.shape == (..., S, E) assert vxs.shape == exs.shape[: -1], "Sequence embedding dimensions do not match." lengths = np.sum(vxs != 0, -1).flatten() # (X,) seqs = F.reshape(exs, (-1, ) + exs.shape[-2:]) # (X, S, E) toembed = [ s[..., :l, :] for s, l in zip(F.separate(seqs, 0), lengths) if l != 0 ] # Y x [(S1, E), (S2, E), ...] hs, ys = birnn(None, toembed) # (2, Y, E), Y x [(S1, 2*E), (S2, 2*E), ...] if return_seqs: ys = F.pad_sequence(ys) # (Y, S, 2*E) ys = F.reshape(ys, ys.shape[:-1] + (2, EMBED)) # (Y, S, 2, E) ys = F.mean(ys, -2) # (Y, S, E) if ys.shape[0] == lengths.size: ys = F.reshape(ys, exs.shape) # (..., S, E) return ys embeds = np.zeros((lengths.size, vxs.shape[-1], EMBED), dtype=np.float32) # (X, S, E) idxs = np.nonzero(lengths) # (Y,) embeds = F.scatter_add(embeds, idxs, ys) # (X, S, E) embeds = F.reshape(embeds, exs.shape) # (..., S, E) return embeds # (..., S, E) hs = F.mean(hs, 0) # (Y, E) if hs.shape[0] == lengths.size: hs = F.reshape(hs, vxs.shape[:-1] + (EMBED, )) # (..., E) return hs # Add zero values back to match original shape embeds = np.zeros((lengths.size, EMBED), dtype=np.float32) # (X, E) idxs = np.nonzero(lengths) # (Y,) embeds = F.scatter_add(embeds, idxs, hs) # (X, E) embeds = F.reshape(embeds, vxs.shape[:-1] + (EMBED, )) # (..., E) return embeds # (..., E)
def __call__(self, atoms_feat, pair_feat, global_feat, atom_idx, pair_idx, start_idx, end_idx): # 1) Pass the Dense layer a_f_d = self.dense_for_atom(atoms_feat) p_f_d = self.dense_for_pair(pair_feat) g_f_d = self.dense_for_global(global_feat) # 2) Update the edge vector start_node = a_f_d[start_idx] end_node = a_f_d[end_idx] g_f_extend_with_pair_idx = g_f_d[pair_idx] concat_p_v = functions.concat((p_f_d, start_node, end_node, g_f_extend_with_pair_idx)) update_p = self.update_for_pair(concat_p_v) # 3) Update the node vector # 1. get sum edge feature of all nodes using scatter_add method zero = self.xp.zeros(a_f_d.shape, dtype=self.xp.float32) sum_edeg_vec = functions.scatter_add(zero, start_idx, update_p) + \ functions.scatter_add(zero, end_idx, update_p) # 2. get degree of all nodes using scatter_add method one = self.xp.ones(p_f_d.shape, dtype=self.xp.float32) degree = functions.scatter_add(zero, start_idx, one) + \ functions.scatter_add(zero, end_idx, one) # 3. get mean edge feature of all nodes mean_edge_vec = sum_edeg_vec / degree # 4. concating g_f_extend_with_atom_idx = g_f_d[atom_idx] concat_a_v = functions.concat((a_f_d, mean_edge_vec, g_f_extend_with_atom_idx)) update_a = self.update_for_atom(concat_a_v) # 4) Update the global vector out_shape = g_f_d.shape ave_p = get_mean_feat(update_p, pair_idx, out_shape, self.xp) ave_a = get_mean_feat(update_a, atom_idx, out_shape, self.xp) concat_g_v = functions.concat((ave_a, ave_p, g_f_d), axis=1) update_g = self.update_for_global(concat_g_v) # 5) Skip connection if self.skip_intermediate: # Skip intermediate feature, used for first layer. new_a_f = update_a + a_f_d new_p_f = update_p + p_f_d new_g_f = update_g + g_f_d else: # Skip input feature, used all layer except first layer. # input feature must be same dimension with updated feature. new_a_f = update_a + atoms_feat new_p_f = update_p + pair_feat new_g_f = update_g + global_feat # 6) dropout if self.dropout_ratio > 0.0: new_a_f = functions.dropout(new_a_f, ratio=self.dropout_ratio) new_p_f = functions.dropout(new_p_f, ratio=self.dropout_ratio) new_g_f = functions.dropout(new_g_f, ratio=self.dropout_ratio) return new_a_f, new_p_f, new_g_f
def get_mean_feat(feat, idx, out_shape, xp): """Return mean node or edge feature in each graph. This method is the same as average pooling about node or edge feature in each graph. """ zero = xp.zeros(out_shape, dtype=xp.float32) sum_vec = functions.scatter_add(zero, idx, feat) one = xp.ones(feat.shape, dtype=xp.float32) degree = functions.scatter_add(zero, idx, one) return sum_vec / degree
def forward(self, xinit, env, train_warm_start_idxs): """ forward function :param xinit: :param env:Pendulum_DX :param train_warm_start_idxs: :return: """ Q = self.xp.zeros((self.n_sc, self.n_sc)) index_diag = self.xp.zeros((self.n_sc, self.n_sc), dtype=bool) self.xp.fill_diagonal(index_diag, True) index_not_diag = self.xp.zeros((self.n_sc, self.n_sc), dtype=bool) index_not_diag[self.xp.tril_indices(self.n_sc, -1)] = True Q = F.scatter_add(Q, self.xp.tril_indices(self.n_sc, -1), self.lower_without_diag) diag_q = F.sigmoid(self.learn_q_logit) Q = F.where(index_diag, diag_q, Q) Q = Q @ Q.T p = self.learn_p # print("p.shape", p.shape) # print(self.n_sc) # p = F.concat((p, self.xp.array(0.0).reshape(1,)), axis=0) # print("Q ", Q) # print("p", p) x_mpc, u_mpc = env.mpc_Q(env.true_dx, xinit, Q, p, u_init=self.xp.transpose(train_warm_start_idxs, axes=(1, 0, 2))) return x_mpc, u_mpc
def seq_rnn_embed(vxs, exs, rnn_layer, initial_state=None, reverse=False): """Embed given sequences using rnn.""" # vxs.shape == (..., S) # exs.shape == (..., S, E) # initial_state == (..., E) assert vxs.shape == exs.shape[: -1], "Sequence embedding dimensions do not match." lengths = np.sum(vxs != 0, -1).flatten() # (X,) seqs = F.reshape(exs, (-1, ) + exs.shape[-2:]) # (X, S, E) if reverse: toembed = [ F.flip(s[..., :l, :], -2) for s, l in zip(F.separate(seqs, 0), lengths) if l != 0 ] # Y x [(S1, E), (S2, E), ...] else: toembed = [ s[..., :l, :] for s, l in zip(F.separate(seqs, 0), lengths) if l != 0 ] # Y x [(S1, E), (S2, E), ...] if initial_state is not None: initial_state = F.reshape(initial_state, (-1, EMBED)) # (X, E) initial_state = initial_state[None, np.flatnonzero(lengths)] # (1, Y, E) hs, ys = rnn_layer(initial_state, toembed) # (1, Y, E), Y x [(S1, 2*E), (S2, 2*E), ...] hs = hs[0] # (Y, E) if hs.shape[0] == lengths.size: hs = F.reshape(hs, vxs.shape[:-1] + (EMBED, )) # (..., E) return hs # Add zero values back to match original shape embeds = np.zeros((lengths.size, EMBED), dtype=np.float32) # (X, E) idxs = np.nonzero(lengths) # (Y,) embeds = F.scatter_add(embeds, idxs, hs) # (X, E) embeds = F.reshape(embeds, vxs.shape[:-1] + (EMBED, )) # (..., E) return embeds # (..., E)
def __call__(self, h, edge_index): # add self node feature new_h = h messages = h[edge_index[0]] new_h = functions.scatter_add(new_h, edge_index[1], messages) # apply MLP new_h = self.mlp(new_h) if self.dropout_ratio > 0.0: new_h = functions.dropout(new_h, ratio=self.dropout_ratio) return new_h
def compute_cos_angle(vert,face,theta,xp): # Obsolite: replaced with compute_cos_angle_sub which is more efficient cos_angle = Variable(xp.cos(theta)) sin_angle = Variable(xp.sin(theta)) for f in face: n = len(f) id_p = xp.array([f[(i-1)%n] for i in range(n+1)]) id = xp.array([f[i%n] for i in range(n+1)]) id_n = xp.array([f[(i+1)%n] for i in range(n)]) L = F.sum((vert[id_p] - vert[id])**2, axis=1) D = F.sum((vert[id_n] - vert[ id_p[:-1] ])**2, axis=1) c1 = (L[:n]+L[1:]-D)/(2*F.sqrt(L[:n]*L[1:])) s1 = F.sqrt(1-c1**2) # trigonometric addition formula c0 = cos_angle[f] s0 = sin_angle[f] cos_angle = F.scatter_add(cos_angle,f,-c0 + c0*c1 - s0*s1) sin_angle = F.scatter_add(sin_angle,f,-s0 + c0*s1 + s0*c1) return cos_angle
def check_forward(self, a_data, b_data): a = chainer.Variable(a_data) b = chainer.Variable(b_data) y = functions.scatter_add(a, self.slices, b) self.assertEqual(y.data.dtype, numpy.float32) # Test to make sure that the input values are not changed numpy.testing.assert_equal(cuda.to_cpu(a.data), self.a_data_original) a_data_copy = cuda.to_cpu(a_data).copy() numpy.add.at(a_data_copy, self.slices, cuda.to_cpu(b_data)) numpy.testing.assert_equal(a_data_copy, cuda.to_cpu(y.data))
def __call__(self, h, batch, h0=None, is_real_node=None): # --- Readout part --- h1 = functions.concat((h, h0), axis=1) if h0 is not None else h g1 = functions.sigmoid(self.i_layer(h1)) g2 = self.activation(self.j_layer(h1)) g = g1 * g2 # sum along node axis y = self.xp.zeros((int(batch[-1]) + 1, self.out_dim), dtype=numpy.float32) y = functions.scatter_add(y, batch, g) y = self.activation_agg(y) if self.concat_n_info: n_nodes = self.xp.zeros(y.shape[0], dtype=self.xp.float32) n_nodes = functions.scatter_add(n_nodes, batch, self.xp.ones(batch.shape[0])) y = functions.concat((y, n_nodes.reshape((-1, 1)))) return y
def compute_curvature(vert,face,xp): # Obsolite: replaced with compute_curvature_sub which is more efficient K = Variable(xp.full((len(vert),), 2*xp.pi)) for f in face: n = len(f) id_p = xp.array([f[(i-1)%n] for i in range(n+1)]) id = xp.array([f[i%n] for i in range(n+1)]) id_n = xp.array([f[(i+1)%n] for i in range(n)]) L = F.sum((vert[id_p] - vert[id])**2, axis=1) D = F.sum((vert[id_n] - vert[ id_p[:-1] ])**2, axis=1) c1 = (L[:n]+L[1:]-D)/(2*F.sqrt(L[:n]*L[1:])) K = F.scatter_add(K,f,-F.arccos(c1)) return K
def __call__(self, h, batch, axis=1, **kwargs): if self.activation is not None: h = self.activation(h) else: h = h if self.mode == 'sum': y = self.xp.zeros((batch[-1] + 1, h.shape[1]), dtype=self.xp.float32) y = functions.scatter_add(y, batch, h) else: raise ValueError('mode {} is not supported'.format(self.mode)) return y
def interpolate_trilinear(grid, lin_ind_frustrum, voxel_coords, img_shape, frustrum_depth): xp = chainer.cuda.get_array_module(voxel_coords) batch, num_feats, height, width, depth = grid.shape lin_ind_frustrum = lin_ind_frustrum.astype( "int32") # indexを指定するだけなので勾配を流す必要はない x_indices = voxel_coords[2, :] y_indices = voxel_coords[1, :] z_indices = voxel_coords[0, :] x0 = x_indices.astype("int32") y0 = y_indices.astype("int32") z0 = z_indices.astype("int32") x1 = x0 + 1 y1 = y0 + 1 z1 = z0 + 1 x1 = xp.clip(x1, 0, width - 1) y1 = xp.clip(y1, 0, height - 1) z1 = xp.clip(z1, 0, depth - 1) x = x_indices - x0 y = y_indices - y0 z = z_indices - z0 # output = torch.zeros(batch, num_feats, img_shape[0]*img_shape[1]*depth).cuda() output = xp.zeros( (batch, num_feats, img_shape[0] * img_shape[1] * frustrum_depth), dtype="float32") added = grid[:, :, x0, y0, z0] * (1 - x) * (1 - y) * (1 - z) + \ grid[:, :, x1, y0, z0] * x * (1 - y) * (1 - z) + \ grid[:, :, x0, y1, z0] * (1 - x) * y * (1 - z) + \ grid[:, :, x0, y0, z1] * (1 - x) * (1 - y) * z + \ grid[:, :, x1, y0, z1] * x * (1 - y) * z + \ grid[:, :, x0, y1, z1] * (1 - x) * y * z + \ grid[:, :, x1, y1, z0] * x * y * (1 - z) + \ grid[:, :, x1, y1, z1] * x * y * z make_slice = MakeSlice() output = F.scatter_add(output, make_slice[:, :, lin_ind_frustrum], added) output = output.reshape(batch, num_feats, frustrum_depth, img_shape[0], img_shape[1]) return output
def forward(self, xinit, env, train_warm_start_idxs): """ forward function :param xinit: :param env:Pendulum_DX :param train_warm_start_idxs: :return: """ Q = self.xp.zeros((self.n_sc, self.n_sc)) index_diag = self.xp.zeros((self.n_sc, self.n_sc), dtype=bool) self.xp.fill_diagonal(index_diag, True) index_not_diag = self.xp.zeros((self.n_sc, self.n_sc), dtype=bool) index_not_diag[self.xp.tril_indices(self.n_sc, -1)] = True Q = F.scatter_add(Q, self.xp.tril_indices(self.n_sc, -1), self.lower_without_diag) diag_q = F.sigmoid(self.learn_q_logit) Q = F.where(index_diag, diag_q, Q) Q = Q @ Q.T p = self.learn_p Q = OBSERVATION_MATRIX.T @ Q @ OBSERVATION_MATRIX p = p @ OBSERVATION_MATRIX x_mpc, u_mpc = env.mpc_Q(env.true_dx, xinit, Q, p, u_init=self.xp.transpose(train_warm_start_idxs, axes=(1, 0, 2))) return x_mpc, u_mpc
def test_too_many_indices(self): with self.assertRaises(type_check.InvalidType): functions.scatter_add(self.a_data, (0, 0, 0, 0), self.b_data)
def test_multiple_ellipsis(self): with self.assertRaises(ValueError): functions.scatter_add( self.a_data, (Ellipsis, Ellipsis), self.b_data)
def test_multiple_ellipsis(self): with self.assertRaises(ValueError): functions.scatter_add(self.a_data, (Ellipsis, Ellipsis), self.b_data)
def f(a, b): y = functions.scatter_add(a, self.slices, b) return y * y
def forward(self, inputs, device): a, b = inputs y = functions.scatter_add(a, self.slices, b) return y,
def f(a, b): return functions.scatter_add(a, self.slices, b)
def test_requires_broadcasting(self): with self.assertRaises(ValueError): functions.scatter_add(self.a_data, slice(0, 2), self.b_data)
def __call__(self, h, edge_index, edge_attr): next_h = self.root_weight(h) features = self.edge_weight(h).reshape(-1, self.n_edge_types, self.out_channels) messages = features[edge_index[0], edge_attr, :] return functions.scatter_add(next_h, edge_index[1], messages)
def dataset_loss(self, data_iter, warmstart=None): """ calculate loss and set warm start this loss is :param loader: :param warmstart: :return: """ true_q, true_p = self.env.true_dx.get_true_obj() losses = [] iter_before = data_iter.epoch while data_iter.epoch < iter_before + 1: next_batch = data_iter.next() next_batch = dataset.concat_examples(next_batch) x_inits = next_batch[0] xs = next_batch[1] us = next_batch[2] idxs = next_batch[3] n_batch = x_inits.shape[0] dx = self.env.true_dx if not self.is_lower_triangle and not self.is_strange_observation: print("not lower triangle and not strange observation") assert type(self.net) == Pendulum_Net_cost_logit q = F.sigmoid(self.net.learn_q_logit) p = F.sqrt(q) * self.net.learn_p _, pred_u = self.env.mpc(self.env.true_dx, x_inits, q, p, u_init=xp.transpose(warmstart[idxs], axes=(1, 0, 2))) pred_u = F.transpose(pred_u, axes=(1, 0, 2)).array warmstart[idxs] = pred_u elif not self.is_strange_observation: print("lower triangle and not strange observation") assert type(self.net) == Pendulum_Net_cost_lower_triangle Q = xp.zeros((self.n_sc, self.n_sc)) index_diag = xp.zeros((self.n_sc, self.n_sc), dtype=bool) xp.fill_diagonal(index_diag, True) index_not_diag = xp.zeros((self.n_sc, self.n_sc), dtype=bool) index_not_diag[xp.tril_indices(self.n_sc, -1)] = True Q = F.scatter_add(Q, xp.tril_indices(self.n_sc, -1), self.net.lower_without_diag) # index_not_diag = xp.zeros((self.n_state, self.n_state), dtype=bool) # index_not_diag[xp.tril_indices(self.n_state, -1)] = True # Q = F.scatter_add(Q, xp.tril_indices(self.n_state, -1), self.net.lower_without_diag) diag_q = F.sigmoid(self.net.learn_q_logit) Q = F.where(index_diag, diag_q, Q) Q = Q @ Q.T p = self.net.learn_p p = F.concat((p, xp.array(0.0).reshape(1, )), axis=0) _, pred_u = self.env.mpc_Q(self.env.true_dx, x_inits, Q, p, u_init=xp.transpose(warmstart[idxs], axes=(1, 0, 2))) pred_u = F.transpose(pred_u, axes=(1, 0, 2)).array warmstart[idxs] = pred_u elif not self.is_lower_triangle: print("not lower triangle and strange observation") assert type( self.net) == Pendulum_Net_cost_logit_strange_obervation q = F.sigmoid(self.net.learn_q_logit) p = F.sqrt(q) * self.net.learn_p # p = self.learn_p Q = xp.zeros((self.n_sc, self.n_sc)) index_diag = xp.zeros((self.n_sc, self.n_sc), dtype=bool) xp.fill_diagonal(index_diag, True) Q = F.where(index_diag, q, Q) Q = OBSERVATION_MATRIX.T @ Q @ OBSERVATION_MATRIX p = p @ OBSERVATION_MATRIX _, pred_u = self.env.mpc_Q(self.env.true_dx, x_inits, Q, p, u_init=xp.transpose(warmstart[idxs], axes=(1, 0, 2))) pred_u = F.transpose(pred_u, axes=(1, 0, 2)).array warmstart[idxs] = pred_u else: print("lower triangle and strange observation") assert type( self.net ) == Pendulum_Net_cost_lower_triangle_strange_obervation Q = xp.zeros((self.n_sc, self.n_sc)) index_diag = xp.zeros((self.n_sc, self.n_sc), dtype=bool) xp.fill_diagonal(index_diag, True) index_not_diag = xp.zeros((self.n_sc, self.n_sc), dtype=bool) index_not_diag[xp.tril_indices(self.n_sc, -1)] = True Q = F.scatter_add(Q, xp.tril_indices(self.n_sc, -1), self.net.lower_without_diag) diag_q = F.sigmoid(self.net.learn_q_logit) Q = F.where(index_diag, diag_q, Q) Q = Q @ Q.T Q = OBSERVATION_MATRIX.T @ Q @ OBSERVATION_MATRIX p = self.net.learn_p p = p @ OBSERVATION_MATRIX _, pred_u = self.env.mpc_Q(self.env.true_dx, x_inits, Q, p, u_init=xp.transpose(warmstart[idxs], axes=(1, 0, 2))) pred_u = F.transpose(pred_u, axes=(1, 0, 2)).array warmstart[idxs] = pred_u assert pred_u.shape == us.shape squared_loss = (us - pred_u) * (us - pred_u) # print(squared_loss.shape) loss = xp.mean(squared_loss) losses.append(loss) loss = xp.stack(losses).mean() return loss