# 阿里巴巴机器翻译团队：将TVM引入TensorFlow中以优化GPU上的神经机器翻译

## 批量Matmul

### 什么是批量matmul？

void BatchedGemm(input A, input B, output C, M, N, K, batch_dimension) {
for (int i = 0; i < batch_dimension; ++i)  {
DoGemm(A[i],B[i],C[i],M,K,N)
}
}

## 批量matmul计算

# computation representation
A = tvm.placeholder((batch, M, K), name='A')
B = tvm.placeholder((batch, K, N), name='B')
k = tvm.reduce_axis((0, K), 'k')
C = tvm.compute((batch, M, N),
lambda b, y, x: tvm.sum(A[b, y, k] * B[b, k, x], axis = k),
name = 'C')

## Schedule优化

### 调整块/线程的参数

  # thread indices
# block partitioning
BB, FF, MM, PP = s[C].op.axis
BBFF = s[C].fuse(BB, FF)
MMPP = s[C].fuse(MM, PP)
s[C].bind(by, block_y)
s[C].bind(bx, block_x)
vty, ty = s[C].split(ty_block, nparts = vthread_y)
vtx, tx = s[C].split(tx_block, nparts = vthread_x)
s[C].reorder(by, bx, vty, vtx, ty, tx)
s[C].reorder(by, bx, ty, tx)
s[C].bind(vtx, thread_xz)

## 将matmul与其他运算融合

# computation representation
A = tvm.placeholder((batch_size, features, M, K), name='A')
# the shape of B is (N, K) other than (K, N) is because B is transposed is this fusion pattern
B = tvm.placeholder((batch_size, features, N, K), name='B')
ENTER = tvm.placeholder((batch_size, 1, M, N), name = 'ENTER')
k = tvm.reduce_axis((0, K), 'k')
C = tvm.compute(
(batch_size, features, M, N),
lambda yb, yf, m, x: tvm.sum(A[yb, yf, m, k] * B[yb, yf, x, k], axis = k),
name = 'C')
D = topi.broadcast_add(C, ENTER)

# computation representation
A = tvm.placeholder((batch_size, features, M, K), name='A')
B = tvm.placeholder((batch_size, features, K, N), name='B')
k = tvm.reduce_axis((0, K), 'k')
C = tvm.compute(
(batch_size, M, features, N),
lambda yb, m, yf, x: tvm.sum(A[yb, yf, m, k] * B[yb, yf, k, x], axis = k),
name = 'C')

### Fusion的性能

· tf-r1.4 BatchMatmul：513.9 us

· tf-r1.4 BatchMatmulTranspose（separate）：541.9 us

· TVM BatchMatmul：37.62us

·TVM BatchMatmulTranspose（fused）：38.39 us

## 参考

|
7天前
|

win10上使用gpu版的tensorflow
win10上使用gpu版的tensorflow
38 0
|
7天前
|

TensorFlow识别GPU难道就这么难吗？还是我的GPU有问题？
TensorFlow识别GPU难道就这么难吗？还是我的GPU有问题？
109 0
|
7天前
|
TensorFlow 算法框架/工具 异构计算
Windows部署TensorFlow后识别GPU失败，原因是啥？
Windows部署TensorFlow后识别GPU失败，原因是啥？
52 0
|
7天前
|
TensorFlow 算法框架/工具 异构计算
TensorFlow检测GPU是否可用
TensorFlow检测GPU是否可用
13 0
|
7天前
|
TensorFlow 算法框架/工具 C++

25 0
|
7天前
|

TensorFlow与GPU加速：提升深度学习性能
【4月更文挑战第17天】本文介绍了TensorFlow如何利用GPU加速深度学习， GPU的并行处理能力适合处理深度学习中的矩阵运算，显著提升性能。TensorFlow通过CUDA和cuDNN库支持GPU，启用GPU只需简单代码。GPU加速能减少训练时间，使训练更大、更复杂的模型成为可能，但也需注意成本、内存限制和编程复杂性。随着技术发展，GPU将继续在深度学习中发挥关键作用，而更高效的硬件解决方案也将备受期待。
70 2
|
7天前
|

【4月更文挑战第17天】本文探讨了如何优化TensorFlow模型的性能，重点介绍了超参数调整和训练技巧。超参数如学习率、批量大小和层数对模型性能至关重要。文章提到了三种超参数调整策略：网格搜索、随机搜索和贝叶斯优化。此外，还分享了训练技巧，包括学习率调度、早停、数据增强和正则化，这些都有助于防止过拟合并提高模型泛化能力。结合这些方法，可构建更高效、健壮的深度学习模型。
34 2
|
7天前
|

31 2
|
7天前
|

63 2
|
7天前
|

Linux Ubuntu配置CPU与GPU版本tensorflow库的方法
Linux Ubuntu配置CPU与GPU版本tensorflow库的方法
91 1