TensorFlow C++ API中SubSlice触发IsAligned()校验失败问题求助

使用TensorFlow C++ API编写自定义神经网络时，将形状为[65000, 784]的张量重塑为[2031, 32, 784]的批量张量，通过SubSlice方法获取每个[32, 784]的批量。当执行Tensor imageBatch(imageBatches.SubSlice(num));（num=1时）触发错误：

2024-04-21 22:27:18.856998: F ./tensorflow/core/framework/tensor.h:888] Check failed: IsAligned() ptr = 0x55555569f130

相关代码片段：

for (int num = 0; num < dataSize; num++) {
            vector<Tensor> output1;
            //auto d1 = DeepCopy(scope, imageBatches.SubSlice(num));
            //auto d2 = DeepCopy(scope, labelBatches.SubSlice(num));
            //TF_CHECK_OK(session->Run({d1, d2}, &output1));
            Tensor imageBatch(imageBatches.SubSlice(num));
            Tensor labelBatch(labelBatches.SubSlice(num));
            //std::cout << "Image batch: " << imageBatch.matrix<float>() << std::endl;
            //std::cout << "imageBatch shape: " << imageBatch.shape().DebugString() << ", labelBatch shape: " << labelBatch.shape().DebugString() <<std::endl;
            TF_CHECK_OK(session->Run({{*features, imageBatch}, {*this->labels, labelBatch}}, apply_gradients, {}, nullptr));
        }

错误原因 #

SubSlice返回的是原张量的内存视图（TensorSlice），它并未独立分配内存，而是直接引用原张量的部分数据。TensorFlow的Tensor对象要求内存地址满足特定对齐规则（比如适配SIMD指令的对齐要求），而SubSlice得到的视图起始地址可能不满足该规则，导致构造新Tensor时触发IsAligned()断言失败。

修复方法 #

• 方法1：使用DeepCopy创建对齐的独立张量
  取消注释代码中的DeepCopy相关逻辑，通过TensorFlow的DeepCopy操作生成内存对齐的独立张量，再传入Run方法：

for (int num = 0; num < dataSize; num++) {
    vector<Tensor> output1;
    auto d1 = DeepCopy(scope, imageBatches.SubSlice(num));
    auto d2 = DeepCopy(scope, labelBatches.SubSlice(num));
    TF_CHECK_OK(session->Run({{*features, d1}, {*this->labels, d2}}, apply_gradients, {}, nullptr));
}

• 方法2：直接传递SubSlice结果，避免构造新Tensor
  session->Run支持直接传入TensorSlice对象，无需手动构造新Tensor，可直接跳过内存对齐检查问题：

for (int num = 0; num < dataSize; num++) {
    vector<Tensor> output1;
    TF_CHECK_OK(session->Run({{*features, imageBatches.SubSlice(num)}, {*this->labels, labelBatches.SubSlice(num)}}, apply_gradients, {}, nullptr));
}

• 方法3：手动复制数据到对齐张量
  如果需要独立的Tensor对象，可以手动创建内存对齐的Tensor，再将SubSlice的数据复制进去：

for (int num = 0; num < dataSize; num++) {
    vector<Tensor> output1;
    // 创建形状为[32,784]的对齐Tensor（标签张量形状请根据实际情况调整）
    Tensor imageBatch(DT_FLOAT, TensorShape({32, 784}));
    Tensor labelBatch(DT_FLOAT, TensorShape({32, /* 标签维度 */}));
    
    // 复制SubSlice的数据到新Tensor
    tensorflow::CopyTensor(imageBatches.SubSlice(num), &imageBatch);
    tensorflow::CopyTensor(labelBatches.SubSlice(num), &labelBatch);
    
    TF_CHECK_OK(session->Run({{*features, imageBatch}, {*this->labels, labelBatch}}, apply_gradients, {}, nullptr));
}

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