[pytorch]实现一个自己个Loss函数

[pytorch]实现一个自己个Loss函数

pytorch本身已经为我们提供了丰富而强大的Loss function接口，详情可见Pytorch的十八个损失函数，这些函数已经可以帮我们解决绝大部分的问题，然而，在具体的实践过程中，我们可能发现还是存在需要自己设计Loss函数的情况，下面笔者就介绍一下如何使用pytorch设计自己的损失函数。

Loss类的实现

具体实践过程中，往往需要把loss function实现一个类，实现之后就可以像正常调用pytorch的损失函数一样初始化和调用。具体来说，需要以下两个类函数：

1. init(): 这个函数需要提供损失函数的一些超参。
2. forward():这个函数具体求最终的loss。

Loss函数示例

以下是笔者在做一个视频分类的项目的时候设计的，原因是我发现在最终过了softmax之后，各个类别的score相差不大，这是限制准确率的一个很大的问题，所以笔者设计了这个损失函数来对各个score相差不大的情况（最高得分score小于0.5）进行惩罚（penalty）。

类的实现

from torch.autograd import Variable
class MylossFunc(nn.Module):
    def __init__(self, deta):
        super(MylossFunc, self).__init__()
        self.deta = deta
        
    
    def forward(self, out, label):
        out = torch.nn.functional.softmax(out, dim=1)
        m = torch.max(out, 1)[0]
        penalty = self.deta * torch.ones(m.size())
        loss = torch.where(m>0.5, m, penalty)
        loss = torch.sum(loss)
        loss = Variable(loss, requires_grad=True)
        return

函数调用

下面是MylossFunc类初始化的细节，具体在使用的时候就和调用pytorch的loss function一样。

loss_function = MylossFunc(10)
if cuda_available:
    loss_function.cuda()

在训练过程中计算loss：

for data in train_loader:  # 对于训练集的每一个batch
        img, label = data
        if cuda_available:
            img = img.cuda()
            label = label.cuda()
 
        out = cnn( img )  # 送进网络进行输出
        #print('out size: {}'.format(out.size()))
        #print('label size: {}'.format(label.size()))
        
        #out = torch.nn.functional.softmax(out, dim=1)
        loss = loss_function( out, label )  # 获得损失
 
        optimizer.zero_grad()  # 梯度归零
        loss.backward()  # 反向传播获得梯度，但是参数还没有更新
        optimizer.step()  # 更新梯度