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
如图所示，在__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_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()
之后，要做的跟2就差不多了，调用就可以了，调用之前只需要
from _ext import roi_pooling
然后写一个类（跟方法2中的一样，forward 和 backward 中调用 roi_pooling 就好）
