mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
24 lines
1002 B
Python
24 lines
1002 B
Python
import torch
|
|
|
|
|
|
def cal_diff(x: torch.Tensor, y: torch.Tensor) -> tuple[float, float, float, float]:
|
|
if torch.equal(x, y):
|
|
return 0.0, 0.0, 0.0, 0.0
|
|
x, y = x.double(), y.double()
|
|
rmse = ((x - y) * (x - y)).mean().sqrt().item()
|
|
cos_diff = 1 - 2 * (x * y).sum().item() / max((x * x + y * y).sum().item(), 1e-12)
|
|
chebyshev_dis = (x - y).abs().max().item()
|
|
kl = ((torch.exp(x - y) - 1) - (x - y)).mean().item()
|
|
return cos_diff, rmse, chebyshev_dis, kl
|
|
|
|
|
|
def assert_close(x: torch.Tensor, y: torch.Tensor, name: str, eps: float = 1e-4) -> None:
|
|
assert x.shape == y.shape, f"{x.shape=}, {y.shape=}"
|
|
assert not x.isnan().any() and not y.isnan().any()
|
|
x, y = x.clone(), y.clone()
|
|
x[x.isfinite().logical_not()] = 0
|
|
y[y.isfinite().logical_not()] = 0
|
|
cos_diff, rmse, chebyshev_dis, kl = cal_diff(x, y)
|
|
print(f"{name}: {cos_diff=}, {rmse=}, {chebyshev_dis=}, {kl=}")
|
|
assert cos_diff < eps, f"{name}: {cos_diff=}, {rmse=}, {chebyshev_dis=}, {kl=}"
|