forked from ccf-ai-infra/GPUCodeForces
42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
class Model(nn.Module):
|
|
"""
|
|
Computes the pairwise Euclidean distance matrix between two sets of vectors.
|
|
This implementation uses broadcasting, which is memory-intensive.
|
|
"""
|
|
def __init__(self):
|
|
super(Model, self).__init__()
|
|
|
|
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Args:
|
|
x (torch.Tensor): A tensor of shape (N, D).
|
|
y (torch.Tensor): A tensor of shape (M, D).
|
|
|
|
Returns:
|
|
torch.Tensor: A tensor of shape (N, M) where Z[i, j] is the distance
|
|
between x[i] and y[j].
|
|
"""
|
|
# x.unsqueeze(1) -> (N, 1, D)
|
|
# y.unsqueeze(0) -> (1, M, D)
|
|
# The result of subtraction is (N, M, D)
|
|
diff = x.unsqueeze(1) - y.unsqueeze(0)
|
|
# Compute L2 norm along the last dimension
|
|
dist_matrix = torch.norm(diff, p=2, dim=-1)
|
|
return dist_matrix
|
|
|
|
# 测试数据配置
|
|
N = 1024 # Number of vectors in the first set
|
|
M = 1024 # Number of vectors in the second set
|
|
D = 512 # Dimensionality of each vector
|
|
|
|
def get_inputs():
|
|
x = torch.randn(N, D)
|
|
y = torch.randn(M, D)
|
|
return [x, y]
|
|
|
|
def get_init_inputs():
|
|
return [] # No special initialization inputs needed
|