[pytorch]实现一个自己个Loss函数
pytorch本身已经为我们提供了丰富而强大的Loss function接口,详情可见Pytorch的十八个损失函数,这些函数已经可以帮我们解决绝大部分的问题,然而,在具体的实践过程中,我们可能发现还是存在需要自己设计Loss函数的情况,下面笔者就介绍一下如何使用pytorch设计自己的损失函数。
Loss类的实现
具体实践过程中,往往需要把loss function实现一个类,实现之后就可以像正常调用pytorch的损失函数一样初始化和调用。具体来说,需要以下两个类函数:
- init(): 这个函数需要提供损失函数的一些超参。
- 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() # 更新梯度