Я написал собственную модель самообслуживания для маркировки последовательностей в кератах.
import keras.layers as ll
from keras import Model
from keras_pos_embd import TrigPosEmbedding
from keras_multi_head import MultiHeadAttention
inputs = ll.Input(shape=(None,))
x = ll.Embedding(10000, 1024)(inputs)
x = TrigPosEmbedding(mode='add')(x)
x = MultiHeadAttention(head_num=8)(x)
x = ll.Dense(units = 512, activation='relu')(x)
x = ll.Dense(units = 4, activation='softmax')(x)
outputs = x
model = Model(inputs, outputs)
model.summary()
Однако эта модель слишком требовательна к памяти; Я пытаюсь тренироваться на очень длинных последовательностях (длина 20000), и при попытке тренировать это дает мне OOM при попытке выделить тензор с формой [16,20000,20000] (Согласно моим расчетам, выделяя только этот тензор потребует> 150 ГБ ОЗУ!).
Я попытался уменьшить batch_size до 1, но он все равно отказывается помещаться в памяти.
Мне нужно как-то изменить его, чтобы он поместился в 10 ГБ RAM GPU
Как мне это сделать?