zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-22 13:07:51 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

lint fix

Browse Source
This commit is contained in:
NripeshN 2023-12-31 01:36:48 +05:30
parent 15307c5367
commit e7114f4b91
2 changed files with 38 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
python/mlx/nn/losses.py
View File

@ -372,12 +372,13 @@ def log_cosh_loss(
return _reduce(loss, reduction)
def focal_loss(
inputs: mx.array,
targets: mx.array,
alpha: float = 0.25,
gamma: float = 2.0,
reduction: str = "none"
reduction: str = "none",
) -> mx.array:
r"""
Computes the Focal Loss between inputs and targets, which is designed to address
@ -412,11 +413,9 @@ def focal_loss(
return _reduce(focal_loss, reduction)
def dice_loss(
inputs: mx.array,
targets: mx.array,
epsilon: float = 1e-6,
reduction: str = "none"
inputs: mx.array, targets: mx.array, epsilon: float = 1e-6, reduction: str = "none"
) -> mx.array:
r"""
Computes the Dice Loss, which is a measure of overlap between two samples.
@ -438,16 +437,14 @@ def dice_loss(
"""
intersection = mx.sum(inputs * targets, axis=1)
cardinality = mx.sum(inputs + targets, axis=1)
dice_score = (2. * intersection + epsilon) / (cardinality + epsilon)
dice_score = (2.0 * intersection + epsilon) / (cardinality + epsilon)
loss = 1 - dice_score
return _reduce(loss, reduction)
def iou_loss(
inputs: mx.array,
targets: mx.array,
epsilon: float = 1e-6,
reduction: str = "none"
inputs: mx.array, targets: mx.array, epsilon: float = 1e-6, reduction: str = "none"
) -> mx.array:
r"""
Computes the Intersection over Union (IoU) Loss, which is a measure of the
@ -474,13 +471,14 @@ def iou_loss(
return _reduce(loss, reduction)
def contrastive_loss(
anchors: mx.array,
positives: mx.array,
negatives: mx.array,
margin: float = 1.0,
p: int = 2,
reduction: str = "none"
reduction: str = "none",
) -> mx.array:
r"""
Computes the Contrastive Loss for a set of anchor, positive, and negative samples.
@ -507,13 +505,14 @@ def contrastive_loss(
return _reduce(loss, reduction)
def tversky_loss(
inputs: mx.array,
targets: mx.array,
alpha: float = 0.5,
beta: float = 0.5,
epsilon: float = 1e-6,
reduction: str = "none"
reduction: str = "none",
) -> mx.array:
r"""
Computes the Tversky Loss, a generalization of the Dice Loss, allowing more control over false
@ -538,7 +537,9 @@ def tversky_loss(
intersection = mx.sum(inputs * targets, axis=1)
false_negatives = mx.sum(inputs * (1 - targets), axis=1)
false_positives = mx.sum((1 - inputs) * targets, axis=1)
tversky_index = (intersection + epsilon) / (intersection + alpha * false_negatives + beta * false_positives + epsilon)
tversky_index = (intersection + epsilon) / (
intersection + alpha * false_negatives + beta * false_positives + epsilon
)
loss = 1 - tversky_index
return _reduce(loss, reduction)

18
python/tests/test_nn.py
View File

@ -801,7 +801,9 @@ class TestNN(mlx_tests.MLXTestCase):
pt = mx.exp(-ce_loss)
expected_loss = -alpha * ((1 - pt) ** gamma) * ce_loss
expected_loss = mx.mean(expected_loss)
loss = nn.losses.focal_loss(inputs, targets, alpha=alpha, gamma=gamma, reduction="mean")
loss = nn.losses.focal_loss(
inputs, targets, alpha=alpha, gamma=gamma, reduction="mean"
)
self.assertAlmostEqual(loss.item(), expected_loss.item(), places=6)
def test_dice_loss(self):
@ -810,7 +812,7 @@ class TestNN(mlx_tests.MLXTestCase):
epsilon = 1e-6
intersection = mx.sum(inputs * targets, axis=1)
cardinality = mx.sum(inputs + targets, axis=1)
dice_score = (2. * intersection + epsilon) / (cardinality + epsilon)
dice_score = (2.0 * intersection + epsilon) / (cardinality + epsilon)
expected_loss = 1 - dice_score
expected_loss = mx.mean(expected_loss)
loss = nn.losses.dice_loss(inputs, targets, epsilon=epsilon, reduction="mean")
@ -838,7 +840,9 @@ class TestNN(mlx_tests.MLXTestCase):
negative_distance = mx.sqrt(mx.power(anchors - negatives, p).sum(axis=1))
expected_loss = mx.maximum(positive_distance - negative_distance + margin, 0)
expected_loss = mx.mean(expected_loss)
loss = nn.losses.contrastive_loss(anchors, positives, negatives, margin=margin, p=p, reduction="mean")
loss = nn.losses.contrastive_loss(
anchors, positives, negatives, margin=margin, p=p, reduction="mean"
)
self.assertAlmostEqual(loss.item(), expected_loss.item(), places=6)
def test_tversky_loss(self):
@ -850,10 +854,14 @@ class TestNN(mlx_tests.MLXTestCase):
intersection = mx.sum(inputs * targets, axis=1)
false_negatives = mx.sum(inputs * (1 - targets), axis=1)
false_positives = mx.sum((1 - inputs) * targets, axis=1)
tversky_index = (intersection + epsilon) / (intersection + alpha * false_negatives + beta * false_positives + epsilon)
tversky_index = (intersection + epsilon) / (
intersection + alpha * false_negatives + beta * false_positives + epsilon
)
expected_loss = 1 - tversky_index
expected_loss = mx.mean(expected_loss)
loss = nn.losses.tversky_loss(inputs, targets, alpha=alpha, beta=beta, epsilon=epsilon, reduction="mean")
loss = nn.losses.tversky_loss(
inputs, targets, alpha=alpha, beta=beta, epsilon=epsilon, reduction="mean"
)
self.assertAlmostEqual(loss.item(), expected_loss.item(), places=6)