(self, coords)
| 56 | return int(2 * 3 * self.frequency_num) |
| 57 | |
| 58 | def forward(self, coords): |
| 59 | device = coords.device |
| 60 | dtype = coords.dtype |
| 61 | N = coords.size(0) |
| 62 | |
| 63 | # (batch_size, num_context_pt, coord_dim) |
| 64 | coords_mat = np.asarray(coords.cpu()) |
| 65 | batch_size = coords_mat.shape[0] |
| 66 | num_context_pt = coords_mat.shape[1] |
| 67 | |
| 68 | # compute the dot product between [deltaX, deltaY] and each unit_vec |
| 69 | # (batch_size, num_context_pt, 1) |
| 70 | angle_mat1 = np.expand_dims(np.matmul(coords_mat, self.unit_vec1), axis=-1) |
| 71 | # (batch_size, num_context_pt, 1) |
| 72 | angle_mat2 = np.expand_dims(np.matmul(coords_mat, self.unit_vec2), axis=-1) |
| 73 | # (batch_size, num_context_pt, 1) |
| 74 | angle_mat3 = np.expand_dims(np.matmul(coords_mat, self.unit_vec3), axis=-1) |
| 75 | |
| 76 | # (batch_size, num_context_pt, 6) |
| 77 | angle_mat = np.concatenate([angle_mat1, angle_mat1, angle_mat2, angle_mat2, angle_mat3, angle_mat3], axis=-1) |
| 78 | # (batch_size, num_context_pt, 1, 6) |
| 79 | angle_mat = np.expand_dims(angle_mat, axis=-2) |
| 80 | # (batch_size, num_context_pt, frequency_num, 6) |
| 81 | angle_mat = np.repeat(angle_mat, self.frequency_num, axis=-2) |
| 82 | # (batch_size, num_context_pt, frequency_num, 6) |
| 83 | angle_mat = angle_mat * self.freq_mat |
| 84 | # (batch_size, num_context_pt, frequency_num*6) |
| 85 | spr_embeds = np.reshape(angle_mat, (batch_size, num_context_pt, -1)) |
| 86 | |
| 87 | # make sinuniod function |
| 88 | # sin for 2i, cos for 2i+1 |
| 89 | # spr_embeds: (batch_size, num_context_pt, frequency_num*6=input_embed_dim) |
| 90 | spr_embeds[:, :, 0::2] = np.sin(spr_embeds[:, :, 0::2]) # dim 2i |
| 91 | spr_embeds[:, :, 1::2] = np.cos(spr_embeds[:, :, 1::2]) # dim 2i+1 |
| 92 | |
| 93 | return torch.from_numpy(spr_embeds.reshape(N,-1)).to(dtype).to(device) |
| 94 | |
| 95 | |
| 96 |
nothing calls this directly
no outgoing calls
no test coverage detected