Block-sparse GPU kernels

OpenAI Blog 工具
gpu-kernels block-sparse optimization tensorflow open-source

摘要

OpenAI 发布 block-sparse GPU kernels，这是一款用于在 GPU 上进行高效稀疏矩阵乘法的工具，可以减少神经网络操作的计算量和内存占用。

我们发布了针对一类被忽视的神经网络架构的高度优化 GPU kernels：具有块稀疏权重的网络。根据所选的稀疏性，这些 kernels 的运行速度可以比 cuBLAS 或 cuSPARSE 快几个数量级。我们已经用它们在文本情感分析、文本和图像生成建模中取得了最先进的结果。
查看原文
查看缓存全文

缓存时间: 2026/04/20 14:55

# 块稀疏 GPU 核函数 来源：https://openai.com/index/block-sparse-gpu-kernels/ ```python 1from blocksparse.matmul import BlocksparseMatMul 2import tensorflow as tf 3import numpy as np 4 5hidden_size = 4096 6block_size = 32 7minibatch_size = 64 8 9# 创建（随机）稀疏性模式 10sparsity = np.random.randint(2, size=(hidden_size//block_size,hidden_size//block_size)) 11 12# 初始化稀疏矩阵乘法对象 13bsmm = BlocksparseMatMul(sparsity, block_size=block_size) 14 15# 图的输入 16x = tf.placeholder(tf.float32, shape=[None, hidden_size]) 17 18# 初始化块稀疏权重 19w = tf.get_variable("w", bsmm.w_shape, dtype=tf.float32) 20 21# 块稀疏矩阵乘法 22y = bsmm(x, w) 23 24# 运行 25sess = tf.InteractiveSession() 26sess.run(tf.global_variables_initializer()) 27result = sess.run([y], feed_dict = {x: np.ones((minibatch_size,hidden_size), dtype='float32')}) 28print(result) ```

相似文章