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

申高峰

2023-12-01

总说

使用方式:
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选项

>>> LANG = MATMUL_LANG
>>> matmul = tc.define(lang, name="matmul")
>>> mat1, mat2 = torch.randn(3, 4).cuda(), torch.randn(4, 5).cuda()
>>> out = matmul(mat1, mat2, options=Options("mlp"))

补充

我们要处理不同size的tensor呢?

如果有新的size输入, TC后端会重新编译定义, 所以你啥都不用做. 当然如果size一直不变, 那么TC编译一次.

import tensor_comprehensions as tc
import torch
lang = """
def matmul(float(M,N) A, float(N,K) B) -> (output) {
  output(i, j) +=! A(i, kk) * B(kk, j)
}
"""
matmul = tc.define(lang, name="matmul")
mat1, mat2 = torch.randn(3, 4).cuda(), torch.randn(4, 5).cuda()
out1 = matmul(mat1, mat2)

# different input sizes
# 重新编译了
mat3, mat4 = torch.randn(100, 400).cuda(), torch.randn(400, 500).cuda()
out2 = matmul(mat3, mat4)

一个字符串中可以放多个TC定义

如果有多个定义的话, 我们不用写多个字符串, 全部写入一个即可, 比如:

import tensor_comprehensions as tc
import torch
lang = """
def matmul(float(M,N) A, float(N,K) B) -> (output) {
  output(i, j) +=! A(i, kk) * B(kk, j)
}
def abs(float(M, N) A) -> (O1) {
  O1(m, n) = fabs(A(m, n))
}
"""
#指定`name`即可
matmul = tc.define(lang, name="matmul")
mat1, mat2 = torch.randn(3, 4).cuda(), torch.randn(4, 5).cuda()
out = matmul(mat1, mat2)

#指定`name`即可
abs = tc.define(lang, name="abs")
A = torch.randn(3, 4).cuda()
out = abs(A)

注入已有的cuda代码

如果我们已经有了高效的cuda代码, 想不用以前非常麻烦的各种方法, 这里的一个比较简单的方式就是, 直接用TC来进行

import tensor_comprehensions as tc
import torch
lang = """
def add(float(N) A, float(N) B) -> (output) {
    output(i) = A(i) + B(i) + 1
}
"""

# cuda代码
cuda_code = """
extern "C"{
__global__ void my_add(float* __restrict__ output, const float* __restrict__ A, const float* __restrict B)
{
    int t = threadIdx.x;
    output[t] = A[t] + B[t];
}
}
"""

add = tc.define(lang, name="add", inject_kernel="my_add", cuda_code=cuda_code)
a, b = torch.randn(100).cuda(), torch.randn(100).cuda()
out = add(a, b, grid=[1, 1, 1], block=[100, 1, 1])

具有scalar的层(需要重写)

有一些层具有scalar, 比如 pooling或是conv, 都有scalar的参数.
因为pytorch现在0.4版本还没正式出来, 所以scalar也看成一个tensor的性质还没完全弄好. 因此, 现在暂时是:

import tensor_comprehensions as tc
import torch
lang = """
def avgpool(float(B, C, H, W) input) -> (output) {{
    output(b, c, h, w) += input(b, c, h * {sH} + kh, w * {sW} + kw) where kh in 0:{kH}, kw in 0:{kW}
}}
"""
#暂时通过`constants`来进行, 以后
avgpool = tc.define(lang, name="avgpool", constants={"sH":1, "sW":1, "kH":2, "kW":2})
inp = torch.ones(32, 3, 10, 10).cuda()
out = avgpool(inp)

可直接调用内置函数

这些内置函数是指cuda内置的函数,

import tensor_comprehensions as tc
import torch
LANG = """
def relu(float(B,M) I) -> (O1){
  O1(b, m) = fmax(I(b, m), 0)
}
"""
relu = tc.define(LANG, name="relu")
inp = torch.randn(100, 128).cuda()
out = relu(inp)

支持的函数在:
http://docs.nvidia.com/cuda/cuda-math-api/group__CUDA__MATH__SINGLE.html#group__CUDA__MATH__SINGLE

直接调用database的layers

对于一些常见的层, 我们可以直接调用现有的字符串

import tensor_comprehensions as tc
import torch
matmul = tc.define(tc.database['matmul']['lang'], name='matmul')
mat1, mat2 = torch.randn(3, 4).cuda(), torch.randn(4, 5).cuda()
out = matmul(mat1, mat2)

一个pytorch例子

import torch
import tensor_comprehensions as tc
import torch.nn.functional as F
from torch.autograd import Variable

lang_maxpool = """
def maxpool(float(B, C, H, W) input) -> (output) {{
    output(b, c, h, w) max=! input(b, c, h * {sH} + kh, w * {sW} + kw) where kh in 0:{kH}, kw in 0:{kW}
}}
"""


maxpool = tc.define(lang_maxpool, name='maxpool', constants={"sH": 2, "sW": 2, "kH": 2, "kW": 2})
input = torch.randn(1, 4, 8, 8).cuda()
output = maxpool(input)
output_ref = F.max_pool2d(Variable(input), 2)

print('size of input: {}'.format(str(input.size())))
print('input: ', input)
print((output_ref - output).abs())

或是

input = torch.randn(1, 4, 8, 8).cuda()
#即使是用Variable包起来, tc定义的操作仍旧可以进行操作
input_v = Variable(input)
output = maxpool(input_v)
output_ref = F.max_pool2d(Variable(input), 2)