TensorFlow: количество элементов было больше, чем можно представить 32-битным типом вывода

Question

Для большой формы тензора, такой как (72, 7007313, 5), и использования большой сети, такой как эта:

import keras
import numpy as np
from keras.optimizers import Adam
from keras.models import Sequential
from keras.layers import Dense, Flatten

Y = np.random.randint(low=0, high=1, dtype=np.int16, size=(72, 2))
X = np.random.randint(low=1, high=5, dtype=np.int16, size=(72, 7007313, 5))

model = keras.models.Sequential()
model.add(Dense(1024, input_shape=X.shape[1:], activation='relu'))
model.add(Dense(1024, activation='relu'))
model.add(Dense(1024, activation='relu'))
model.add(Dense(1024, activation='relu'))
model.add(Dense(1024, activation='relu'))
model.add(Dense(1024, activation='relu'))
model.add(Flatten())
model.add(Dense(2, activation='softmax'))
model.compile(optimizer=Adam(),loss='binary_crossentropy', metrics=['accuracy'])
model.summary()
model.fit(X, Y, epochs=1, batch_size=1)

Обучение этого тензора с использованием этой нейронной сети с batch_size = 1 дает следующую ошибку :

_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense_1 (Dense)              (None, 7007313, 1024)     6144      
_________________________________________________________________
dense_2 (Dense)              (None, 7007313, 1024)     1049600   
_________________________________________________________________
dense_3 (Dense)              (None, 7007313, 1024)     1049600   
_________________________________________________________________
dense_4 (Dense)              (None, 7007313, 1024)     1049600   
_________________________________________________________________
dense_5 (Dense)              (None, 7007313, 1024)     1049600   
_________________________________________________________________
dense_6 (Dense)              (None, 7007313, 1024)     1049600   
_________________________________________________________________
flatten_1 (Flatten)          (None, 7175488512)        0         
_________________________________________________________________
dense_7 (Dense)              (None, 2)                 1435097702
=================================================================
Total params: 14,356,231,170
Trainable params: 14,356,231,170
Non-trainable params: 0
_________________________________________________________________
Epoch 1/1
Traceback (most recent call last):
tensorflow.python.framework.errors_impl.InvalidArgumentError:
Number of elements was larger than representable by 32-bit output type

Это связано с тем, что в сети имеется 14 356 231 170 обучаемых элементов, что превышает стандартную настройку TensorFlow, которая индексирует тензоры с использованием int32 (т.е. не более 4 294 967 296 обучаемых элементов).

Итак, мой вопрос: Как я могу изменить TensorFlow для индексирования тензоров, используя вместо этого int64 ?

Answer 1

OP_REQUIRES(
          ctx, FastBoundsCheck(size, std::numeric_limits<int32>::max()),
          errors::InvalidArgument("Number of elements was larger than "
                                  "representable by 32-bit output type"));

Согласно дизайну , вы не можете иметь TensorFlow, чтобы иметь размеры тензор выше предела int32.

Ссылка на исходный код