在GPU内存中加载TensorFlow模型

时间:2018-09-04 18:15:04

标签: python python-3.x tensorflow gpu

我写了一段简单的代码来自动化一个相对繁琐的任务,即图像分类。问题是我从互联网上获取一张图像,将TF模型加载到GPU中,对其进行处理,然后转到第二张图像。这导致大量的内存和时间浪费。您建议的是将模型加载到GPU内存本身的最佳方法,这样我就不会再三再三地加载相同的模型了。如果无法在GPU内存上加载模型,则可以将CPU加载到内存中,这样可以节省时间。

该模型是一个相对较小的模型,我严重怀疑是否存在任何内存问题。

PS-我不是技术人员,也不完全了解这些概念。我正在尝试破解它,因此,如果我在这里没有任何意义,请原谅我。

以下是我正在使用的代码:-

   from keras.models import load_model
from helpers import resize_to_fit
from imutils import paths
import numpy as np
import imutils
import cv2
import pickle
import sys
from copy import copy
from helpers import diff_pixels
relative_dir = sys.argv[1]
MODEL_FILENAME = relative_dir + "\\htext_model.hdf5"
MODEL_LABELS_FILENAME = relative_dir + "\\model_labels.dat"


# Load up the model labels (so we can translate model predictions to actual letters)
with open(MODEL_LABELS_FILENAME, "rb") as f:
    lb = pickle.load(f)

# Load the trained neural network
model = load_model(MODEL_FILENAME)
image_file = sys.argv[2]
image = cv2.imread(image_file, cv2.IMREAD_GRAYSCALE)
cropped_image = image[5:72,5:242]
image = cv2.resize(cropped_image,(237,67))
# threshold the image (convert it to pure black and white)
thresh = cv2.threshold(image,127,255,cv2.THRESH_BINARY)[1]

letter_positions = []
x_pos = 0
letter_flag = 0
for i in range(0,236):
    black_pixels = 67 - cv2.countNonZero(thresh[:,i:i+1])
    # print(black_pixels)
    if black_pixels < 2:
        if letter_flag == 1 and i-x_pos >= 8:
            letter_positions.append([x_pos-1,i])
            letter_flag = 0
    else:
        if letter_flag == 0:
            x_pos = i
            letter_flag = 1

predictions = []
for i in range(0,len(letter_positions)):
    # Grab the coordinates of the letter in the image
    x1,x2 = letter_positions[i]
    if x1<0: x1=0
    letter_image = image[:,x1:x2]
    # Re-size the letter image to 20x20 pixels to match training data
    letter_image = resize_to_fit(letter_image, 40, 40)

    # Turn the single image into a 4d list of images to make Keras happy
    letter_image = np.expand_dims(letter_image, axis=2)
    letter_image = np.expand_dims(letter_image, axis=0)

    # Ask the neural network to make a prediction
    prediction = model.predict(letter_image)

    # Convert the one-hot-encoded prediction back to a normals letter
    letter = lb.inverse_transform(prediction)[0]
    predictions.append(letter)

# Print the htext's text
htext_text = "".join(predictions)
print(format(htext_text))

0 个答案:

没有答案