Skip to content

Instantly share code, notes, and snippets.

@ilivans
Created March 30, 2018 15:58
Show Gist options
  • Save ilivans/fb2d61d9b5bc3d82d3d0e6eb04cf4778 to your computer and use it in GitHub Desktop.
Save ilivans/fb2d61d9b5bc3d82d3d0e6eb04cf4778 to your computer and use it in GitHub Desktop.
from __future__ import print_function
import os
import gc
from time import sleep
import keras.backend as K
import numpy as np
import psutil
from keras.models import Sequential
from keras.layers import Dense, Activation
# set parameters:
train_size = 512
test_size = 128
num_features = 400
batch_size = 32
hidden_size = 1024 * 8
epochs = 1
print("Generate data...")
np.random.seed(42)
x_train = np.random.rand(train_size, num_features)
y_train = np.random.randint(0, 1, size=train_size)
x_test = np.random.rand(test_size, num_features)
y_test = np.random.randint(0, 1, size=test_size )
print('x_train shape:', x_train.shape)
print('x_test shape:', x_test.shape)
print('Build model...')
model = Sequential()
model.add(Dense(hidden_size, input_shape=(num_features,)))
model.add(Dense(1))
model.add(Activation('sigmoid'))
model.compile(loss='binary_crossentropy',
optimizer='adam',
metrics=['accuracy'])
# print("Train...")
# model.fit(x_train, y_train,
# batch_size=batch_size,
# epochs=epochs)
print("Predict...")
K.set_learning_phase(0)
process = psutil.Process(os.getpid())
while True:
print("rss={}MB\tvms={}MB".format(
process.memory_info().rss // 10 ** 6,
process.memory_info().vms // 10 ** 6
))
sleep(1)
y_pred = model.predict(x_test, batch_size=256)
del y_pred
gc.collect()
@bzamecnik
Copy link

In Keras==2.1.6, tensorflow-gpu==1.8.0, on GTX 1070 I observed constant memory usage (rss=1108MB vms=26964MB). But I'm still hunting a leak in our prediction code.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment