Pytorch学习 (十二) ----- Pytorch中如何使用Tensor Comprehensions

总说

使用方式:
1. 定义TC语言, 然后传入 tc.define
2. 创建torch tensors(用不用Variable包起来都不影响, pytorch0.4版本都没必要加Variable, 直接tensor即可)
3. 跑这一层就行

import tensor_comprehensions as tc
import torch
MATMUL_LANG = """
def matmul(float(M,N) A, float(N,K) B) -> (output) {
  output(i, j) +=! A(i, kk) * B(kk, j)
}
"""
# `name`必须是定义中的`matmul`
matmul = tc.define(MATMUL_LANG, name="matmul")
mat1, mat2 = torch.randn(3, 4).cuda(), torch.randn(4, 5).cuda()
out = matmul(mat1, mat2)

这里看到tc.define

tc.define

这个是针对我们写的TC语言进行定义, 常用选项有:
lang: 表达式
name: 操作名称
backward: 如果我们这一层需要进行反向传播, 那么就要写这个反向传播的表达式!
training: 这个只在lang中有2个TC定义, 这2个分别代表前向的定义和反向的定义时才起作用. True就是开启训练模式.
constants:针对我们的表达式中有常数, 比如 针对3*3的conv kernel, 或是针对2*2的pooling之类的.
inject_kernel,cuda_code在后面注入已有的代码有讲到.

tc.TcUnit

tc.define返回的是TC definition, 也就是tensor_comprehensions.TcUint类.
matmul = tc.define(MATMUL_LANG, name="matmul")
我们拿到这个类的实例matmul后, 它就是一个操作了, 我们可以传入参数进行计算.
主要有三个参数
input没啥好说的
options: 这个是核的计算映射方式, 主要有五种
— pointwise: 如果这个核类似于pointwise操作(比如, tensor的pointwise乘法啊)
— mlp: 如果核类似与线性操作
— conv:如果核类似与卷积操作(比如 MaxPooling之类的)
— group_conv: 核类似与卷积操作(加了分组而已)
— naive: 这是默认的(效果一般不会好)
当然, 你可能会说,我都不知道哪个选项好, 都试试就行. 上一篇博客中写了自动调参, 看运行时间即可.
cache没啥好说的
此时我们将上面的代码加上options选项

>>&g