222 lines
6.9 KiB
Python
222 lines
6.9 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
# Copyright (C) 2026 David (Corey) Herrera
|
|
"""
|
|
Comparison tests for loss functions.
|
|
"""
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from compare_utils import ComparisonTest, TestRunner
|
|
|
|
|
|
class TestMSELoss(ComparisonTest):
|
|
def __init__(self):
|
|
super().__init__("mse_loss", tolerance=1e-6)
|
|
|
|
def generate_input(self, seed=42):
|
|
np.random.seed(seed)
|
|
pred = np.random.randn(32, 10).astype(np.float32)
|
|
target = np.random.randn(32, 10).astype(np.float32)
|
|
return (pred, target)
|
|
|
|
def pytorch_forward(self, pred, target):
|
|
return F.mse_loss(pred, target, reduction='mean')
|
|
|
|
def foundry_params(self):
|
|
return {"reduction": "mean"}
|
|
|
|
|
|
class TestL1Loss(ComparisonTest):
|
|
def __init__(self):
|
|
super().__init__("l1_loss", tolerance=1e-6)
|
|
|
|
def generate_input(self, seed=42):
|
|
np.random.seed(seed)
|
|
pred = np.random.randn(32, 10).astype(np.float32)
|
|
target = np.random.randn(32, 10).astype(np.float32)
|
|
return (pred, target)
|
|
|
|
def pytorch_forward(self, pred, target):
|
|
return F.l1_loss(pred, target, reduction='mean')
|
|
|
|
def foundry_params(self):
|
|
return {"reduction": "mean"}
|
|
|
|
|
|
class TestSmoothL1Loss(ComparisonTest):
|
|
def __init__(self):
|
|
super().__init__("smooth_l1_loss", tolerance=1e-5)
|
|
|
|
def generate_input(self, seed=42):
|
|
np.random.seed(seed)
|
|
pred = np.random.randn(32, 10).astype(np.float32)
|
|
target = np.random.randn(32, 10).astype(np.float32)
|
|
return (pred, target)
|
|
|
|
def pytorch_forward(self, pred, target):
|
|
return F.smooth_l1_loss(pred, target, reduction='mean', beta=1.0)
|
|
|
|
def foundry_params(self):
|
|
return {"reduction": "mean", "beta": 1.0}
|
|
|
|
|
|
class TestCrossEntropyLoss(ComparisonTest):
|
|
def __init__(self):
|
|
super().__init__("cross_entropy", tolerance=1e-5)
|
|
|
|
def generate_input(self, seed=42):
|
|
np.random.seed(seed)
|
|
# Logits (not probabilities)
|
|
logits = np.random.randn(32, 10).astype(np.float32)
|
|
# Class indices
|
|
targets = np.random.randint(0, 10, 32).astype(np.int64)
|
|
return (logits, targets)
|
|
|
|
def pytorch_forward(self, logits, targets):
|
|
return F.cross_entropy(logits, targets, reduction='mean')
|
|
|
|
def foundry_params(self):
|
|
return {"reduction": "mean"}
|
|
|
|
|
|
class TestNLLLoss(ComparisonTest):
|
|
def __init__(self):
|
|
super().__init__("nll_loss", tolerance=1e-5)
|
|
|
|
def generate_input(self, seed=42):
|
|
np.random.seed(seed)
|
|
# Log probabilities
|
|
logits = np.random.randn(32, 10).astype(np.float32)
|
|
log_probs = logits - np.log(np.sum(np.exp(logits), axis=1, keepdims=True))
|
|
# Class indices
|
|
targets = np.random.randint(0, 10, 32).astype(np.int64)
|
|
return (log_probs.astype(np.float32), targets)
|
|
|
|
def pytorch_forward(self, log_probs, targets):
|
|
return F.nll_loss(log_probs, targets, reduction='mean')
|
|
|
|
def foundry_params(self):
|
|
return {"reduction": "mean"}
|
|
|
|
|
|
class TestBinaryCrossEntropyWithLogits(ComparisonTest):
|
|
def __init__(self):
|
|
super().__init__("binary_cross_entropy_with_logits", tolerance=1e-5)
|
|
|
|
def generate_input(self, seed=42):
|
|
np.random.seed(seed)
|
|
logits = np.random.randn(32, 10).astype(np.float32)
|
|
# Binary targets
|
|
targets = np.random.randint(0, 2, (32, 10)).astype(np.float32)
|
|
return (logits, targets)
|
|
|
|
def pytorch_forward(self, logits, targets):
|
|
return F.binary_cross_entropy_with_logits(logits, targets, reduction='mean')
|
|
|
|
def foundry_params(self):
|
|
return {"reduction": "mean"}
|
|
|
|
|
|
class TestKLDivLoss(ComparisonTest):
|
|
def __init__(self):
|
|
super().__init__("kl_div", tolerance=1e-5)
|
|
|
|
def generate_input(self, seed=42):
|
|
np.random.seed(seed)
|
|
# Log probabilities for input
|
|
logits = np.random.randn(32, 10).astype(np.float32)
|
|
log_probs = logits - np.log(np.sum(np.exp(logits), axis=1, keepdims=True))
|
|
# Probabilities for target
|
|
target_logits = np.random.randn(32, 10).astype(np.float32)
|
|
target_probs = np.exp(target_logits) / np.sum(np.exp(target_logits), axis=1, keepdims=True)
|
|
return (log_probs.astype(np.float32), target_probs.astype(np.float32))
|
|
|
|
def pytorch_forward(self, log_probs, target):
|
|
return F.kl_div(log_probs, target, reduction='batchmean')
|
|
|
|
def foundry_params(self):
|
|
return {"reduction": "batchmean"}
|
|
|
|
|
|
class TestHuberLoss(ComparisonTest):
|
|
def __init__(self):
|
|
super().__init__("huber_loss", tolerance=1e-5)
|
|
|
|
def generate_input(self, seed=42):
|
|
np.random.seed(seed)
|
|
pred = np.random.randn(32, 10).astype(np.float32)
|
|
target = np.random.randn(32, 10).astype(np.float32)
|
|
return (pred, target)
|
|
|
|
def pytorch_forward(self, pred, target):
|
|
return F.huber_loss(pred, target, reduction='mean', delta=1.0)
|
|
|
|
def foundry_params(self):
|
|
return {"reduction": "mean", "delta": 1.0}
|
|
|
|
|
|
class TestCosineEmbeddingLoss(ComparisonTest):
|
|
def __init__(self):
|
|
super().__init__("cosine_embedding_loss", tolerance=1e-5)
|
|
|
|
def generate_input(self, seed=42):
|
|
np.random.seed(seed)
|
|
x1 = np.random.randn(32, 64).astype(np.float32)
|
|
x2 = np.random.randn(32, 64).astype(np.float32)
|
|
# Labels: 1 for similar, -1 for dissimilar
|
|
y = np.random.choice([-1, 1], 32).astype(np.int64)
|
|
return (x1, x2, y)
|
|
|
|
def pytorch_forward(self, x1, x2, y):
|
|
return F.cosine_embedding_loss(x1, x2, torch.tensor(y), reduction='mean')
|
|
|
|
def foundry_params(self):
|
|
return {"reduction": "mean", "margin": 0.0}
|
|
|
|
|
|
class TestTripletMarginLoss(ComparisonTest):
|
|
def __init__(self):
|
|
super().__init__("triplet_margin_loss", tolerance=1e-5)
|
|
|
|
def generate_input(self, seed=42):
|
|
np.random.seed(seed)
|
|
anchor = np.random.randn(32, 64).astype(np.float32)
|
|
positive = anchor + np.random.randn(32, 64).astype(np.float32) * 0.1
|
|
negative = np.random.randn(32, 64).astype(np.float32)
|
|
return (anchor, positive.astype(np.float32), negative)
|
|
|
|
def pytorch_forward(self, anchor, positive, negative):
|
|
return F.triplet_margin_loss(anchor, positive, negative, margin=1.0, reduction='mean')
|
|
|
|
def foundry_params(self):
|
|
return {"reduction": "mean", "margin": 1.0}
|
|
|
|
|
|
def get_all_tests():
|
|
"""Return all loss function tests."""
|
|
return [
|
|
TestMSELoss(),
|
|
TestL1Loss(),
|
|
TestSmoothL1Loss(),
|
|
TestCrossEntropyLoss(),
|
|
TestNLLLoss(),
|
|
TestBinaryCrossEntropyWithLogits(),
|
|
TestKLDivLoss(),
|
|
TestHuberLoss(),
|
|
TestCosineEmbeddingLoss(),
|
|
TestTripletMarginLoss(),
|
|
]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
debug = "--debug" in sys.argv
|
|
|
|
runner = TestRunner()
|
|
for test in get_all_tests():
|
|
runner.add_test(test)
|
|
|
|
runner.run_all(debug=debug)
|
|
runner.save_results()
|