【DeepLearning】【PyTorch 】PyTorch 如何自定义损失函数（Loss Function）？

转自：https://www.zhihu.com/question/66988664

1. 直接利用torch.Tensor提供的接口

因为只是需要自定义loss，而loss可以看做对一个或多个Tensor的混合计算，比如计算一个三元组的Loss(Triplet Loss)，我们只需要如下操作：(假设输入的三个(anchor, positive, negative)张量维度是 batch_size * 400<即triplet(net的输出)>)

import torch
import torch.nn as nn
import torch.nn.functional as func
class TripletLossFunc(nn.Module):def __init__(self, t1, t2, beta):super(TripletLossFunc, self).__init__()self.t1 = t1self.t2 = t2self.beta = betareturndef forward(self, anchor, positive, negative):    matched = torch.pow(func.pairwise_distance(anchor, positive), 2)(func.pairwise_distance(anchor, positive), 2)part_1 = torch.clamp(matched - mismatched, min=self.t1)part_2 = torch.clamp(matched, min=self.t2)dist_hinge = part_1 + self.beta * part_2loss = torch.mean(dist_hinge)return loss

如图所示，在__init__()中定义超参数，在forward()中定义计算过程就可以了，全程使用torch提供的张量计算接口(道理上同样可以使用numpy和scipy的，不过感觉效率会低一点)，该方法可调用cuda(仅限仅使用了torch接口或者python内建方法)，(即你可以直接使用实例化对象的.cuda()方法)

因为继承了nn.Module，所以这个Loss类在实例化之后可以直接运行__call__()方法，也就是

a = TripletLossFunc(...)
loss = a(anchor, positive, negative)

就可以了。这是第一种方法。

2. 利用PyTorch的numpy/scipy扩展

如果你细心的话你会注意到我在上面使用了torch.nn.functional模块的函数，那么，问题来了，万一需要的计算不在这个模块中怎么办？
那么，问题来了，万一需要的计算不在这个模块中怎么办？
官网教程在此 (官网教程是自定义一个快速傅里叶变换在网络中，我们也可以定义操作然后用在loss中)
https://pytorch.org/tutorials/advanced/numpy_extensions_tutorial.html
你需要做的操作其实只多了一步：

import torch
from torch.autograd import Function
from troch.autograd import Variable
class OwnOp(Function):def forward(input_tensor):tensor = input_tensor.numpy()...... # 其它 numpy/scipy 操作result = ......return torch.Tensor(result)def backward(grad_output):

注意，你只需要定义 forward() 和 backward() 两个方法就可以了，务必需要先调用输入的 .numpy() 方法，返回需要把返回值变成 torch.Tensor。

写到这里，基本满足大部分需求了，但是，有了另外一个问题，如果我需要计算的东西很多(比如需要涉及到像素级别的计算)或者很复杂，或者numpy/scipy中没有这些操作怎么办？

恩，那就只有最后一种方法了，不过需要你有一定的C语言基础和会使用CUDA编程(据传MSRA很多写CUDA很熟练的神)

3. 写一个PyTorch的C扩展

恩。。。。最近再被这个玩意折腾，还在学cuda23333，对于这个，我先给个官网的教程

PyTorch C扩展
https://pytorch.org/tutorials/advanced/cpp_extension.html#

以及某大神写的一个roi_pooling的C扩展 ROI

具体的话，需要你先定义最基本的C/CUDA运算

/* triplest_cal.c */
#include <TH/TH.h>
#include <math.h>
int triplet_cal_forward(...)
{
    // 计算代码
}
int triplet_cal_backward(...)
{
    // 计算代码
}

/* triplet_cal.h *? int triplet_cal_forward(...); int triplet_cal_backward(...); 

注意们这里的文件名必须跟模块名相同，比如你的模块名是 triplet_cal，呢文件名就如上。
然后 forward，backward那两个函数名也必须遵照这个格式。
因为 PyTorch 自己有一个 Parse 用来解析头文件，从而进行相关的运算

cuda 同理，也需要定义 triplet_cal_cuda.c 和 triplet_cal_cuda.h
cuda 需要额外定义 cuda 运算

/* triplet_cal_kernel.cu */
#ifdef __cplusplus
extern "C"{
    
#endif#include <stdio.h>
#include <math.h>
#include <float.h>
#include "triplet_cal_kernel.h"
}
/* 我还不会CUDA233333 */

然后，你需要定义 build.py，用来注册这个扩展，使它被 PyTorch 接受(我自己的扩展还没写到这一步，所以我把roi_pooling的拿过来了23333，这个模块名就叫做roi_pooling)

import os
import torch
from torch.utils.ffi import create_extensionsources = ['src/roi_pooling.c']
headers = ['src/roi_pooling.h']
definex = []
with_cuda = Falseif torch.cuda.is_available():print('Including CUDA code.')sources += ['src/roi_pooling_cuda.c']headers += ['src/roi_pooling.cuda.h']defines += [('WITH_CUDA', None)]with_cuda = Truethis_file = os.path.dirname(os.path.realpath(__file__))
print(this_file)
extra_objects = ['src/cuda/roi_pooling.cu.o']
extra_objects = [os.path.join(this_file, name) for fname in extra_objects]ffi = create_extension('_ext.roi_pooling',headers=headers,sources=sources,define_macros=defines,relative_to = __file__,with_cuda=with_cudaextra_objects=extra_objects
)if __name__ == '__main__'ffi.build()

之后，要做的跟2就差不多了，调用就可以了，调用之前只需要

from _ext import roi_pooling

然后写一个类（跟方法2中的一样，forward 和 backward 中调用 roi_pooling 就好）