tensorflow-gpu1.7.0运行CollaboNet训练报embedding_lookup索引越界问题

问题背景 #

为复现目标开源项目提出的模型，按照依赖要求通过pip命令安装指定版本依赖：pip install tensorflow-gpu==1.7.0。安装完成后环境初始化流程无异常，启动训练任务时触发运行时错误。

完整报错堆栈 #

====NCBI_MODEL Training=====
Traceback (most recent call last):
  File "run.py", line 226, in <module>
    sess, batch_idx, infoInput=intOuts, tbWriter=tbWriter)
  File "/home/youssef/Documents/collabonet/CollaboNet-master/model/RunModel.py", line 82, in train1epoch
    feed_dict=feed_dict1)
  File "/home/youssef/.local/lib/python2.7/site-packages/tensorflow/python/client/session.py", line 905, in run
    run_metadata_ptr)
  File "/home/youssef/.local/lib/python2.7/site-packages/tensorflow/python/client/session.py", line 1140, in _run
    feed_dict_tensor, options, run_metadata)
  File "/home/youssef/.local/lib/python2.7/site-packages/tensorflow/python/client/session.py", line 1321, in _do_run
    run_metadata)
  File "/home/youssef/.local/lib/python2.7/site-packages/tensorflow/python/client/session.py", line 1340, in _do_call
    raise type(e)(node_def, op, message)
tensorflow.python.framework.errors_impl.InvalidArgumentError: indices[2,15] = 1651 is not in [0, 2)
     [[Node: BC5CDR-diseaseWE/embedding_lookup = Gather[Tindices=DT_INT32, Tparams=DT_FLOAT, _class=["loc:@BC5CDR-diseaseWE/embed"], validate_indices=true, _device="/job:localhost/replica:0/task:0/device:CPU:0"](BC5CDR-diseaseWE/embed, _arg_X_0_0)]]

Caused by op u'BC5CDR-diseaseWE/embedding_lookup', defined at:
  File "run.py", line 119, in <module>
    modelDict[dataSet]['WE']=modelClass.we(args=modelDict[dataSet]['args'])
  File "/home/youssef/Documents/collabonet/CollaboNet-master/model/CollaboNet.py", line 104, in we
    reuse=False)
  File "/home/youssef/Documents/collabonet/CollaboNet-master/ops/embeddingOps.py", line 76, in embedding_lookup
    inputs_embed = tf.nn.embedding_lookup(embedding_table, inputs)
  File "/home/youssef/.local/lib/python2.7/site-packages/tensorflow/python/ops/embedding_ops.py", line 327, in embedding_lookup
    transform_fn=None)
  File "/home/youssef/.local/lib/python2.7/site-packages/tensorflow/python/ops/embedding_ops.py", line 151, in _embedding_lookup_and_transform
    result = _clip(_gather(params[0], ids, name=name), ids, max_norm)
  File "/home/youssef/.local/lib/python2.7/site-packages/tensorflow/python/ops/embedding_ops.py", line 55, in _gather
    return array_ops.gather(params, ids, name=name)
  File "/home/youssef/.local/lib/python2.7/site-packages/tensorflow/python/ops/array_ops.py", line 2698, in gather
    params, indices, validate_indices=validate_indices, name=name)
  File "/home/youssef/.local/lib/python2.7/site-packages/tensorflow/python/ops/gen_array_ops.py", line 2672, in gather
    validate_indices=validate_indices, name=name)
  File "/home/youssef/.local/lib/python2.7/site-packages/tensorflow/python/framework/op_def_library.py", line 787, in _apply_op_helper
    op_def=op_def)
  File "/home/youssef/.local/lib/python2.7/site-packages/tensorflow/python/framework/ops.py", line 3290, in create_op
    op_def=op_def)
  File "/home/youssef/.local/lib/python2.7/site-packages/tensorflow/python/framework/ops.py", line 1654, in __init__
    self._traceback = self._graph._extract_stack()  # pylint: disable=protected-access

InvalidArgumentError (see above for traceback): indices[2,15] = 1651 is not in [0, 2)
     [[Node: BC5CDR-diseaseWE/embedding_lookup = Gather[Tindices=DT_INT32, Tparams=DT_FLOAT, _class=["loc:@BC5CDR-diseaseWE/embed"], validate_indices=true, _device="/job:localhost/replica:0/task:0/device:CPU:0"](BC5CDR-diseaseWE/embed, _arg_X_0_0)]]

已完成的排查操作 #

• 执行import tensorflow as tf、tf.__version__校验依赖版本，返回结果为'1.7.0'，符合项目要求的版本标准
• 怀疑环境误安装CPU版本TensorFlow，执行完全卸载后重新安装指定GPU版本，问题仍复现
• 拉取TensorFlow-GPU 1.7.0官方Docker镜像搭建干净运行环境，执行相同训练流程，报错完全一致

问题根因与修复方案 #

该报错与TensorFlow版本、GPU/CPU运行环境无任何关联，错误核心逻辑可直接从报错信息中定位：

执行embedding_lookup查找操作时，传入的索引张量中存在值为1651的元素，但当前绑定的嵌入表第一维度仅为2，合法索引范围仅为[0, 2)（仅允许取值0、1），触发索引越界校验报错。
结合项目代码逻辑，触发问题的原因是多任务训练架构下的参数传递错位，可按以下步骤修复：

1. 定位嵌入层初始化逻辑
   打开项目路径下model/CollaboNet.py第104行附近的词嵌入构建代码，可发现原代码在构建BC5CDR-disease子任务的嵌入层时，错误将下游二分类输出层的标签数（值为2）作为嵌入表的大小参数传入。正确的参数应为预处理阶段生成的对应数据集的总词表长度（与输入索引张量的最大值1651匹配，即嵌入表大小应不小于1652）。
2. 添加临时校验逻辑确认维度错配
   在ops/embeddingOps.py第76行的tf.nn.embedding_lookup调用前，添加临时校验代码确认参数值：

   print("传入嵌入层的输入张量最大值：", tf.reduce_max(inputs).eval(session=sess))
   print("当前嵌入表第一维度大小：", embedding_table.shape[0])
   

   运行后可确认：传入的输入张量为经过词表映射的单词ID张量（最大值1651，对应总词表规模约1652），而嵌入表大小被错误设置为下游二分类的标签数2，二者维度完全不匹配。
3. 修正多任务参数映射关系
   该项目为多任务联合训练架构，原代码在循环构建不同子任务的嵌入层时，错误将其他任务的输出层参数传入嵌入层构建函数。修正参数传递逻辑，保证每个子任务的嵌入层大小绑定对应数据集预处理阶段生成的词表总长度，即可解决越界报错。

注意事项 #

• 反复重装TensorFlow、切换GPU/CPU版本、更换Docker镜像的操作，完全无法解决维度不匹配的逻辑层错误，无需在这类操作上重复尝试
• 不要通过给tf.nn.embedding_lookup传入validate_indices=False参数临时关闭索引校验，这类操作只会绕过报错检查，导致嵌入输出出现未定义值，最终训练结果完全失真

内容的提问来源于stack exchange，提问作者Josephus