android - Android NN中关于conv2D的问题：输出形状限制

Question

当我尝试 android nn 代码时，在 Conv2D.cpp 中找到了这些：

tflite::Dims<4> im2colDim;                                                    \
im2colDim.sizes[3] = (int)getSizeOfDimension(outputShape, 0);                 \
im2colDim.sizes[2] = (int)getSizeOfDimension(outputShape, 1);                 \
im2colDim.sizes[1] = (int)getSizeOfDimension(outputShape, 2);                 \
im2colDim.sizes[0] = (int)inDepth * filterHeight * filterWidth;               \
                                                                              \
im2colDim.strides[0] = 1;                                                     \
for (int i = 1; i < 4; i++) {                                                 \
    im2colDim.strides[i] = im2colDim.strides[i - 1] * im2colDim.sizes[i - 1]; \
}                                                                             \
                                                                              \
Type* im2colData = nullptr;                                                   \
uint64_t im2colByteSize = sizeof(Type);                                       \
std::unique_ptr<Type[]> im2colGuard;                                          \
for (int i = 0; i < 4; i++) {                                                 \
    im2colByteSize *= im2colDim.sizes[i];                                     \
}                                                                             \
/* http://b/77982879, tflite::optimized_ops::Conv uses int for offsets */     \
if (im2colByteSize >= 0x7fffffff) {                                           \
    LOG(ERROR) << "Conv size is too large, not enough memory";                \
    return false;                                                             \
}                                                                             \
if (im2colByteSize <= kStaticBufferSize) {                                    \
    im2colData = reinterpret_cast<Type*>(static_scratch_buffer);              \
} else {                                                                      \
    im2colData = new (std::nothrow) Type[im2colByteSize / sizeof(Type)];      \
    if (im2colData == nullptr) {                                              \
        LOG(ERROR) << "Conv size is too large, not enough memory";            \
        return false;                                                         \
    }                                                                         \
    im2colGuard.reset(im2colData);                                            \
}

我想知道为什么必须检查 byteSize？

代码链接： https ://android.googlesource.com/platform/frameworks/ml/+/refs/heads/master/nn/common/operations/Conv2D.cpp

163号线附近

多谢你们

Answer 1

以字节为单位的张量总大小 (im2colByteSize) 的两项检查是为了检查：

• 缓冲区的大小高于可以用signed int表示的MAX正数（这是为了防止tflite::optimized_ops::Conv使用signed int表示偏移引起的SEGV）

• 大小足够小，可以避免分配内存并使用 static_scratch_buffer