Keras立即使用所有GPU内存

时间:2019-01-30 20:39:09

标签: python-3.x tensorflow keras

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

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

我尝试做的第一件事是降低网络的复杂性(减少嵌入空间和LSTM) 在尝试使用“ fit_generator”和“ train_onbatch”函数并创建了生成器之后,仍在代码中。 接下来,我尝试将批次大小减小到2,尽管这样做确实非常慢(预计每个时期约600个小时) 在此之后,我试图使在更有效的内存中读取文件的方式有效,即通过使用numpy数组vs列表等。 最后,我尝试使用非gpu版本的TensorFlow,这样做时一分钟内就用完了我所有的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()

1 个答案:

答案 0 :(得分:2)

您可以使用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的问题,您仍然需要选择正确(不是很大)的批处理大小和缓冲区大小。