focal loss代码pytorch
时间: 2023-10-30 08:21:31 浏览: 168
focal loss.py
以下是 PyTorch 中实现 Focal Loss 的代码:
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
class FocalLoss(nn.Module):
def __init__(self, gamma=2, alpha=None, size_average=True):
super(FocalLoss, self).__init__()
self.gamma = gamma
self.alpha = alpha
if isinstance(alpha, (float, int)): self.alpha = torch.Tensor([alpha, 1 - alpha])
if isinstance(alpha, list): self.alpha = torch.Tensor(alpha)
self.size_average = size_average
def forward(self, input, target):
if input.dim() > 2:
input = input.view(input.size(0), input.size(1), -1) # N,C,H,W => N,C,H*W
input = input.transpose(1, 2) # N,C,H*W => N,H*W,C
input = input.contiguous().view(-1, input.size(2)) # N,H*W,C => N*H*W,C
target = target.view(-1, 1)
logpt = F.log_softmax(input)
logpt = logpt.gather(1, target)
logpt = logpt.view(-1)
pt = logpt.data.exp()
if self.alpha is not None:
if self.alpha.type() != input.data.type():
self.alpha = self.alpha.type_as(input.data)
at = self.alpha.gather(0, target.data.view(-1))
logpt = logpt * at
loss = -1 * (1 - pt) ** self.gamma * logpt
if self.size_average: return loss.mean()
else: return loss.sum()
```
可以看到,该代码中首先定义了一个 `FocalLoss` 类,该类继承自 PyTorch 中的 `nn.Module` 类,因此我们可以直接使用该类来定义我们的 Focal Loss 模型。
在 `__init__()` 方法中,我们定义了两个超参数 `gamma` 和 `alpha`。其中 `gamma` 的值默认为 2,即 Focal Loss 中的调节因子。`alpha` 表示每个类别的权重,如果 `alpha` 是一个浮点数,则表示正样本的权重,负样本的权重为 1 - `alpha`。如果 `alpha` 是一个列表,则它的长度应该等于类别数,每个元素表示每个类别的权重。
在 `forward()` 方法中,我们首先将输入的 `input` 和 `target` 二者都展平成一维向量,然后计算损失函数。具体而言,我们首先对 `input` 进行 softmax 操作,然后取出对应类别的概率值 `pt`,接着根据 `alpha` 权重计算加权的对数概率值 `logpt`。最后根据 Focal Loss 的公式计算损失,并返回平均值或总和。
阅读全文