ND4J GPU训练时主机内存分配失败问题求助

修复ND4J Compressors加载错误后，使用GPU训练DeepLearning4J模型时触发以下内存相关异常：

• java.lang.RuntimeException: Failed to allocate 4735031021 bytes from HOST memory
• java.lang.RuntimeException: cudaMalloc failed; Bytes: [4735031021]; Error code [2]; DEVICE [0]

该问题仅在GPU环境下出现，CPU运行时无异常。尝试设置启动参数-Xmx28g -Dorg.bytedeco.javacpp.maxbytes=30G未解决问题，错误发生在首次迭代后的model.fit(allTrainingData)调用环节。

原始代码 #

//First: get the dataset using the record reader. CSVRecordReader handles loading/parsing
int numLinesToSkip = 0;
char delimiter = ',';
RecordReader recordReader = new CSVRecordReader(numLinesToSkip,delimiter);
recordReader.initialize(new FileSplit(new File("vector.txt")));

//Second: the RecordReaderDataSetIterator handles conversion to DataSet objects, ready for use in neural network
int labelIndex = Integer.parseInt("5422");
int numClasses = Integer.parseInt("1170");
int batchSize = 4000;

DataSetIterator iterator = new RecordReaderDataSetIterator.Builder(recordReader, batchSize).classification(labelIndex, numClasses).build();

List<DataSet> trainingData = new ArrayList<>();
List<DataSet> testData = new ArrayList<>();

while (iterator.hasNext()) {
    DataSet allData = iterator.next();
    allData.shuffle();
    SplitTestAndTrain testAndTrain = allData.splitTestAndTrain(0.9);  // Use 90% of data for training
    trainingData.add(testAndTrain.getTrain());
    testData.add(testAndTrain.getTest());
}

DataSet allTrainingData = DataSet.merge(trainingData);
DataSet allTestData = DataSet.merge(testData);

//We need to normalize our data. We'll use NormalizeStandardize (which gives us mean 0, unit variance):       
DataNormalization normalizer = new NormalizerStandardize();
normalizer.fit(allTrainingData);           // Collect the statistics (mean/stdev) from the training data. This does not modify the input data
normalizer.transform(allTrainingData);     // Apply normalization to the training data
normalizer.transform(allTestData);         // Apply normalization to the test data. This is using statistics calculated from the *training* set

long seed = 6;
int firstHiddenLayerSize = labelIndex/6;
int secondHiddenLayerSize = firstHiddenLayerSize/4;

MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
        .seed(seed)
        .activation(Activation.TANH)
        .weightInit(WeightInit.XAVIER)
        //.updater(new Sgd(0.1))
        .updater(Adam.builder().build())
        .l2(1e-4)
        .list()
        .layer(new DenseLayer.Builder().nIn(labelIndex).nOut(firstHiddenLayerSize)
                .build())
        .layer(new DenseLayer.Builder().nIn(firstHiddenLayerSize).nOut(secondHiddenLayerSize)
                .build())
        .layer( new OutputLayer.Builder(LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD)
                .activation(Activation.SOFTMAX) //Override the global TANH activation with softmax for this layer
                .nIn(secondHiddenLayerSize).nOut(numClasses).build())
        .build();

//run the model
MultiLayerNetwork model = new MultiLayerNetwork(conf);
model.init();

//record score once every 100 iterations
model.setListeners(new ScoreIterationListener(100));

for(int i=0; i<5000; i++) {
    model.fit(allTrainingData);
}

//evaluate the model on the test set
Evaluation eval = new Evaluation(numClasses);

INDArray output = model.output(allTestData.getFeatures());

eval.eval(allTestData.getLabels(), output);
log.info(eval.stats());

// Save the Model
File locationToSave = new File("trained-model.zip");
model.save(locationToSave, true);

// Save DataNormalization
NormalizerSerializer ns = NormalizerSerializer.getDefault();
ns.write(normalizer, new File("trained-normalizer.bin"));

更新后的分批训练代码 #

...
DataSetIterator iterator = new RecordReaderDataSetIterator.Builder(recordReader, batchSize).classification(labelIndex, numClasses).build();

List<DataSet> trainingData = new ArrayList<>();

while (iterator.hasNext()) {
    trainingData.add(iterator.next());
}

DataSet allTrainingData = DataSet.merge(trainingData);

// We need to normalize our data. We'll use NormalizeStandardize (which gives us mean 0, unit variance):       
// 此处代码与原始代码一致

// MultiLayerConfiguration conf... 
// 此处代码与原始代码一致
MultiLayerNetwork model = new MultiLayerNetwork(conf);
model.init();

List<DataSet> allTrainingDataBatched = allTrainingData.batchBy(10000);
for (int i=0; i<5000; i++) {
    for (DataSet dataSet: allTrainingDataBatched) {
        model.fit(dataSet);
    }
}
...

修复思路 #

原始代码中每次迭代直接传入合并后的完整训练集，导致GPU或主机内存需要一次性分配大量空间存储整个数据集。修复方案通过batchBy(10000)将完整训练集拆分为更小的批次，每次仅将单个批次的数据传入model.fit()，大幅降低单次训练的内存占用需求，从而避免内存分配失败的异常。

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