tensorflow 在加载大型的embedding模型参数时,会遇到cannot be larger than 2GB
这种问题是,对于每一个变量 variable 由于是基于protobuf存在这大小限制(2G),这个时候,我们需要将embedding拆开,拆分成N等分,来使得每一个
variable都在2G以下;
1 # !/usr/bin/env/python 2 # coding=utf-8 3 import tensorflow as tf 4 import numpy as np 5 6 input_ids = tf.placeholder(dtype=tf.int32, shape=[None,None]) 7 8 num_shards = 3 9 weights = [] 10 weights_shape = np.arange(27).reshape(9, 3) 11 # assert weights_shape[0] % num_shards == 0 12 num_shards_len = (weights_shape.shape[0]) / num_shards 13 assert (weights_shape.shape[0]) % num_shards ==0 14 begin_ = 0 15 ends_ = num_shards_len 16 for i in range(0, num_shards): 17 if (i + 1) * num_shards_len < weights_shape.shape[0]: 18 begin_ = i * num_shards_len 19 if i + 1 == num_shards: 20 ends_ = weights_shape.shape[0] 21 else: 22 ends_ = (i + 1) * num_shards_len 23 else: 24 begin_ = i * num_shards_len 25 ends_ = weights_shape.shape[0] 26 weights_i = tf.get_variable("words-%02d" % i, 27 initializer=tf.constant(weights_shape[begin_: ends_, ])) 28 weights.append(weights_i) 29 30 input_embedding = tf.nn.embedding_lookup(weights, input_ids,partition_strategy="div") 31 32 sess = tf.InteractiveSession() 33 sess.run(tf.global_variables_initializer()) 34 print(sess.run(weights)) 35 36 print(sess.run(input_embedding, feed_dict={input_ids: [[1, 2], [3, 0], [8, 2], [5, 1]]}))
结果为:
[array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]), array([[ 9, 10, 11], [12, 13, 14], [15, 16, 17]]), array([[18, 19, 20], [21, 22, 23], [24, 25, 26]])] [[[ 3 4 5] [ 6 7 8]] [[ 9 10 11] [ 0 1 2]] [[24 25 26] [ 6 7 8]] [[15 16 17] [ 3 4 5]]]
编程是一种快乐,享受代码带给我的乐趣!!!
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】凌霞软件回馈社区,博客园 & 1Panel & Halo 联合会员上线
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】博客园社区专享云产品让利特惠,阿里云新客6.5折上折
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
· [.NET]调用本地 Deepseek 模型
· 一个费力不讨好的项目,让我损失了近一半的绩效!
· .NET Core 托管堆内存泄露/CPU异常的常见思路
· PostgreSQL 和 SQL Server 在统计信息维护中的关键差异
· C++代码改造为UTF-8编码问题的总结
· 一个费力不讨好的项目,让我损失了近一半的绩效!
· 清华大学推出第四讲使用 DeepSeek + DeepResearch 让科研像聊天一样简单!
· 实操Deepseek接入个人知识库
· CSnakes vs Python.NET:高效嵌入与灵活互通的跨语言方案对比
· Plotly.NET 一个为 .NET 打造的强大开源交互式图表库