tf.split(),tf.strided_slice(), tf.slice()

博客介绍了TensorFlow中用于矩阵切割的函数,如tf.split()、tf.slice()和tf.strided_slice()。详细说明了各函数的用法、切割方式及特点,还对比了tf.slice()和tf.strided_slice()在end索引开闭区间上的差异,并给出了使用示例。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

tf.split(): 有两种切割方式
https://blog.youkuaiyun.com/SangrealLilith/article/details/80272346
tf.slice():
https://www.jianshu.com/p/71e6ef6c121b
tf.strided_slice函数:
https://blog.youkuaiyun.com/banana1006034246/article/details/75092388 note:该贴很赞
https://blog.youkuaiyun.com/a13602955218/article/details/81254684

tf.split(): 分成几瓣,每一瓣为一个列表,并共同组成一个list返回.
tf.strided_slice():

tf.stride_slice(data, begin, end)
tf.slice(data, begin, end)
两者都是用于切割矩阵,不同点在于slice的end索引是闭区间,stride_slice的end索引是开区间,所以一个截掉最后一列的小技巧是用stride_slice(data, [0, 0], [rows, -1]),但是如果是用slice(data, [0, 0], [rows, -1])则仍是原矩阵。
注意begin和end不是一个常数,而是一个数组,
例如对于矩阵[1,2,3,4,5],截取前三个,应该用tf.slice(data,[0],[2])
对于
[[1,2,3],
[4,5,6],
[7,8,9]]
截取1,2,4,5,应该用tf.slice(data,[0,0],[1,1])

tf.slice():切出一片
1,函数原型 tf.slice(inputs,begin,size,name=’’)
2,用途:从inputs中抽取部分内容

 inputs:可以是list,array,tensor

 begin:n维列表,begin[i] 表示从inputs中第i维抽取数据时,相对0的起始偏移量,也就是从第i维的begin[i]开始抽取数据

 size:n维列表,size[i]表示要抽取的第i维元素的数目

 有几个关系式如下:

     (1) i in [0,n]

     (2)tf.shape(inputs)[0]=len(begin)=len(size)

     (3)begin[i]>=0   抽取第i维元素的起始位置要大于等于0

     (4)begin[i]+size[i]<=tf.shape(inputs)[i]

tf.split(dimension, num_split, input):dimension的意思就是输入张量的哪一个维度,如果是0就表示对第0维度进行切割。num_split就是切割的数量,如果是2就表示输入张量被切成2份,每一份是一个列表。

将其中的缓存机制去掉,代码变成了 class MLA(layers.Layer): def __init__(self, args: ModelArgs): super().__init__() self.dim = args.dim self.n_heads = args.n_heads self.q_lora_rank = args.q_lora_rank self.kv_lora_rank = args.kv_lora_rank self.qk_nope_head_dim = args.qk_nope_head_dim self.qk_rope_head_dim = args.qk_rope_head_dim self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim self.v_head_dim = args.v_head_dim # 初始化投影层 if self.q_lora_rank == 0: self.wq = layers.Dense(self.n_heads * self.qk_head_dim) else: self.wq_a = layers.Dense(self.q_lora_rank) self.q_norm = RMSNorm(self.q_lora_rank) self.wq_b = layers.Dense(self.n_heads * self.qk_head_dim) self.wkv_a = layers.Dense(self.kv_lora_rank + self.qk_rope_head_dim) self.kv_norm = RMSNorm(self.kv_lora_rank) self.wkv_b = layers.Dense(self.n_heads * (self.qk_nope_head_dim + self.v_head_dim)) self.wo = layers.Dense(self.dim) self.softmax_scale = self.qk_head_dim ** -0.5 if args.max_seq_len > args.original_seq_len: mscale = 0.1 * args.mscale * math.log(args.rope_factor) + 1.0 self.softmax_scale *= mscale * mscale def call(self, x, start_pos, freqs_cis, mask=None): bsz = tf.shape(x)[0] seqlen = tf.shape(x)[1] end_pos = start_pos + seqlen # 查询投影 if self.q_lora_rank == 0: q = self.wq(x) else: q = self.wq_b(self.q_norm(self.wq_a(x))) q = tf.reshape(q, [bsz, seqlen, self.n_heads, self.qk_head_dim]) q_nope, q_pe = tf.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], axis=-1) q_pe = apply_rotary_emb(q_pe, freqs_cis) # 键值投影 kv = self.wkv_a(x) kv, k_pe = tf.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], axis=-1) k_pe = apply_rotary_emb(tf.expand_dims(k_pe, 2), freqs_cis) kv = self.wkv_b(self.kv_norm(kv)) kv = tf.reshape(kv, [bsz, seqlen, self.n_heads, self.qk_nope_head_dim + self.v_head_dim]) k_nope, v = tf.split(kv, [self.qk_nope_head_dim, self.v_head_dim], axis=-1) k = tf.concat([k_nope, tf.tile(k_pe, [1, 1, self.n_heads, 1])], axis=-1) # 注意力计算 q = tf.concat([q_nope, q_pe], axis=-1) scores = tf.einsum("bqhd,bkhd->bhqk", q, k) * self.softmax_scale # 维度调整为qk交互 print(scores.shape) # 此处为(2, 16, 128, 128) if mask is not None: print(mask.shape) # 此处为(128, 128) scores += mask[:, None, :, :] scores += mask[None, None, :, :] scores = tf.nn.softmax(scores, axis=-1) x = tf.einsum("bhqk,bkhd->bqhd", scores, v) # 维度调整 return self.wo(tf.reshape(x, [bsz, seqlen, -1])) 存在以下问题025-03-12 17:01:42.362209: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at strided_slice_op.cc:111 : INVALID_ARGUMENT: Index out of range using input dim 2; input has only 2 dims File "E:\算法模型\DeepSeek-V3-main\inference\model_tf.py", line 248, in call x = x + self.attn(self.attn_norm(x), start_pos, freqs_cis, mask) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "E:\算法模型\DeepSeek-V3-main\inference\model_tf.py", line 159, in call scores += mask[:, None, :, :] ~~~~^^^^^^^^^^^^^^^ tensorflow.python.framework.errors_impl.InvalidArgumentError: Exception encountered when calling layer 'mla' (type MLA). {{function_node __wrapped__StridedSlice_device_/job:localhost/replica:0/task:0/device:CPU:0}} Index out of range using input dim 2; input has only 2 dims [Op:StridedSlice] name: transformer/block/mla/strided_slice/ Call arguments received by layer 'mla' (type MLA): • x=tf.Tensor(shape=(2, 128, 10), dtype=float32) • start_pos=0 • freqs_cis=tf.Tensor(shape=(128, 32), dtype=float32) • mask=tf.Tensor(shape=(128, 128), dtype=float32)
03-13
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值