【DeepLearning】【PyTorch 】PyTorch 如何自定义损失函数（Loss Function）？_pytorch 对输出tensor计算自定义loss

转自：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 = t1
        self.t2 = t2
        self.beta = beta
        return

    def 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_2
        loss = torch.mean(dist_hinge)
        return loss
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

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

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

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

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

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):
1
2
3
4
5
6
7
8
9
10

注意，你只需要定义 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(...)
{
    // 计算代码
}
1
2
3
4
5
6
7
8
9
10
11

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

注意们这里的文件名必须跟模块名相同，比如你的模块名是 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
*/
1
2
3
4
5
6
7
8
9
10
11
12
13

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

import os
import torch
from torch.utils.ffi import create_extension

sources = ['src/roi_pooling.c']
headers = ['src/roi_pooling.h']
definex = []
with_cuda = False

if 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 = True

this_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_cuda
    extra_objects=extra_objects
)

if __name__ == '__main__'
    ffi.build()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

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

from _ext import roi_pooling
1

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