Keras 立即使用所有 GPU 内存

Question

我目前正在尝试对 APK 文件执行静态分析，并拥有大约 50,000 个样本的数据集。为了使它们成为可用于我的网络的格式，每个 APK 都已被反编译并组合 smali 文件，然后组合文件中的所有 smali 命令都被翻译成一个数字，生成的文件是一个 CSV 文件。 然后我尝试使用上述文件作为 Keras 网络的输入，但由于某种原因，我一直遇到 OOM 错误，例如“试图分配 880.21MiB 的内存不足。以下是当前分配摘要。”

我的系统如下： 64GB DDR4 内存 GTX 1080 - 8GB 显存 i5 8600

我尝试的第一件事是降低网络的复杂性（减少嵌入空间和 LSTM） 之后尝试使用“fit_generator”和“train_onbatch”函数并创建生成器 - 仍在代码中。 接下来，我尝试将批量大小减少到 2，尽管这非常慢（预测一个时期大约需要 600 小时） 在此之后，我尝试使文件在内存中的读取方式更加高效，即通过使用 numpy 数组与列表等。 最后，我尝试使用 TensorFlow 的非 GPU 版本，这样做时，我的所有 RAM 64GB 都会在大约一分钟内用完。

我也尝试在 TS 中设置配置设置，即内存增长等，但没有运气

导入操作系统

import numpy as np
import pandas as pd
from keras import Sequential
from keras.layers import Embedding, Conv1D, MaxPooling1D, Dropout, LSTM, Dense
from sklearn.model_selection import train_test_split
from sklearn.utils import shuffle


def read_datasets(path):
    benign = {
        'file_name': np.random.permutation(
            [f for f in os.listdir(os.path.join(path, 'benign')) if
             os.path.isfile(os.path.join(path, 'benign', f))]),
        'label': 0,
        'dir': '/benign'
    }
    malicious = {
        'file_name': np.random.permutation(
            [f for f in os.listdir(os.path.join(path, 'malicious')) if
             os.path.isfile(os.path.join(path, 'malicious', f))]),
        'label': 1,
        'dir': '/malicious'
    }
    b_len = len(benign['file_name'])
    m_len = len(malicious['file_name'])
    result = pd.concat([pd.DataFrame(data=benign, index=[x for x in range(0, b_len)]),
                        pd.DataFrame(data=malicious, index=[x for x in range(b_len + 1, b_len + m_len + 1)])])
    result = shuffle(result)
    result.set_index('file_name', inplace=True)
    return result


def batch_generator(df, batch_size):
    for i in range(0, len(df), batch_size):
        yield preprocess_subset(df[i:i + batch_size]), df['label'][i:i + batch_size]


def get_max_file_len_for_batch(df):
    max_length = float('-inf')
    for row in df.iterrows():
        with open('../../dataset' + os.path.join(row[1]['dir'], 'file_lengths', row[0] + '.length')) as infp:
            x = infp.read()
        if int(x) > max_length:
            max_length = int(x)
    return max_length


def preprocess_subset(df):
    max_file_len = get_max_file_len_for_batch(df)
    X = np.empty((len(df), max_file_len))
    for i, row in enumerate(df.iterrows()):
        data = pd.read_csv('../../dataset' + os.path.join(row[1]['dir'], row[0]), dtype='int16',
                           delimiter=',').values
        if np.max(data) > 256 or np.min(data) < 0:
            print('../../dataset' + os.path.join(row[1]['dir'], row[0]))
        else:
            data = data[data != 0]
            if len(data) > max_file_len:
                max_offset = len(data) - max_file_len
                offset = np.random.randint(max_offset)
                data = data[offset:(max_file_len + offset)]
            else:
                if max_file_len > len(data):
                    max_offset = max_file_len - len(data)
                    offset = np.random.randint(max_offset)
                else:
                    offset = 0
                data = np.pad(data, (offset, max_file_len - len(data) - offset), "constant")
            X[i,] = data
    return X


def model_cnn_lstm():
    max_features = 256
    embedding_size = 50
    model = Sequential()
    model.add(Embedding(max_features, embedding_size))
    model.add(Conv1D(64, 3, padding='valid', activation='relu', strides=1))
    model.add(MaxPooling1D(2, strides=2))
    model.add(Conv1D(64, 3, padding='valid', activation='relu', strides=1))
    model.add(MaxPooling1D(2, strides=2))
    model.add(Dropout(0.5))

    model.add(LSTM(32))
    model.add(Dense(1, activation='sigmoid'))
    model.compile(loss='binary_crossentropy',
                  optimizer='adam',
                  metrics=['accuracy'])
    return model


def run():
    df = read_datasets('../../dataset')

    x_train, x_test, y_train, y_test = train_test_split(df.index, df['label'], test_size=0.33, random_state=42)


    curr_model = model_cnn_lstm()
    x_train = preprocess_subset(df.loc[x_train])

    # for x_batch, y_batch, in batch_generator(x_train, 16):
    curr_model.fit(x_train, y_train, batch_size=16, epochs=5)

    curr_model.save('model.hdf5')


run()

Answer 1

您可以使用 tf.data.Dataset API。 There's a way 仅从您的 csv 文件创建生成器。

dataset = tf.data.experimental.CsvDataset(
    "my_file*.csv",
    [tf.float32,  # Required field, use dtype or empty tensor
     tf.constant([0.0], dtype=tf.float32),  # Optional field, default to 0.0
     tf.int32,  # Required field, use dtype or empty tensor
     ],
    select_cols=[1,2,3]  # Only parse last three columns
)

但这仍然不是最好的方法。如果您将 csv 转换为 TFRecords，训练会更快。Here 是一个很好的帖子。

为了解决 OOM 的问题，您仍然需要选择正确（不大）的批大小和缓冲区大小。