compute the eucilidean distance matrix between embeddings1 and embeddings2 using gpu
(emb1, emb2)
| 134 | |
| 135 | |
| 136 | def pdist_torch(emb1, emb2): |
| 137 | ''' |
| 138 | compute the eucilidean distance matrix between embeddings1 and embeddings2 |
| 139 | using gpu |
| 140 | ''' |
| 141 | m, n = emb1.shape[0], emb2.shape[0] |
| 142 | emb1_pow = torch.pow(emb1, 2).sum(dim=1, keepdim=True).expand(m, n) |
| 143 | emb2_pow = torch.pow(emb2, 2).sum(dim=1, keepdim=True).expand(n, m).t() |
| 144 | dist_mtx = emb1_pow + emb2_pow |
| 145 | dist_mtx = dist_mtx.addmm_(1, -2, emb1, emb2.t()) |
| 146 | # dist_mtx = dist_mtx.clamp(min = 1e-12) |
| 147 | dist_mtx = dist_mtx.clamp(min=1e-12).sqrt() |
| 148 | return dist_mtx |
| 149 | |
| 150 | |
| 151 | def pdist_np(emb1, emb2): |