60 lines
1.8 KiB
Python
60 lines
1.8 KiB
Python
|
from __future__ import absolute_import
|
||
|
|
from __future__ import print_function
|
||
|
|
from __future__ import division
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from torch.nn import functional as F
|
||
|
|
|
||
|
|
|
||
|
|
def compute_distance_matrix(input1, input2, metric='euclidean'):
|
||
|
|
# check input
|
||
|
|
assert isinstance(input1, torch.Tensor)
|
||
|
|
assert isinstance(input2, torch.Tensor)
|
||
|
|
assert input1.dim() == 2, 'Expected 2-D tensor, but got {}-D'.format(input1.dim())
|
||
|
|
assert input2.dim() == 2, 'Expected 2-D tensor, but got {}-D'.format(input2.dim())
|
||
|
|
assert input1.size(1) == input2.size(1)
|
||
|
|
|
||
|
|
if metric == 'euclidean':
|
||
|
|
distmat = euclidean_squared_distance(input1, input2)
|
||
|
|
elif metric == 'cosine':
|
||
|
|
distmat = cosine_distance(input1, input2)
|
||
|
|
else:
|
||
|
|
raise ValueError(
|
||
|
|
'Unknown distance metric: {}. '
|
||
|
|
'Please choose either "euclidean" or "cosine"'.format(metric)
|
||
|
|
)
|
||
|
|
|
||
|
|
return distmat
|
||
|
|
|
||
|
|
|
||
|
|
def euclidean_squared_distance(input1, input2):
|
||
|
|
"""
|
||
|
|
Args:
|
||
|
|
input1 (torch.Tensor): 2-D feature matrix
|
||
|
|
input2 (torch.Tensor): 2-D feature matrix
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
distmat (numpy.ndarray): distance matrix
|
||
|
|
"""
|
||
|
|
m, n = input1.size(0), input2.size(0)
|
||
|
|
distmat = torch.pow(input1, 2).sum(dim=1, keepdim=True).expand(m, n) + \
|
||
|
|
torch.pow(input2, 2).sum(dim=1, keepdim=True).expand(n, m).t()
|
||
|
|
distmat.addmm_(1, -2, input1, input2.t())
|
||
|
|
return distmat.numpy()
|
||
|
|
|
||
|
|
|
||
|
|
def cosine_distance(input1, input2):
|
||
|
|
"""
|
||
|
|
Args:
|
||
|
|
input1 (torch.Tensor): 2-D feature matrix
|
||
|
|
input2 (torch.Tensor): 2-D feature matrix
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
distmat (numpy.ndarray): distance matrix
|
||
|
|
"""
|
||
|
|
input1_normed = F.normalize(input1, p=2, dim=1)
|
||
|
|
input2_normed = F.normalize(input2, p=2, dim=1)
|
||
|
|
distmat = 1 - torch.mm(input1_normed, input2_normed.t())
|
||
|
|
return distmat.numpy()
|