Как бороться с избытком памяти в моем коде keras? - PullRequest
Новая
       45

Как бороться с избытком памяти в моем коде keras?

0 голосов
/ 15 января 2019

Сейчас я работаю над проектом поиска изображений, но запросы изображений и опорных изображений превысили объем памяти графического процессора (24 ГБ) Мой код ниже. 

def infer(queries, db):

        # Query number: 195
        # Reference(DB) number: 1,127
        # Total (query + reference): 1,322

        queries, query_img, references, reference_img = preprocess(queries, db)

        print('test data load queries {} query_img {} references {} reference_img {}'.
              format(len(queries), len(query_img), len(references), len(reference_img)))

        queries = np.asarray(queries)
        query_img = np.asarray(query_img)
        references = np.asarray(references)
        reference_img = np.asarray(reference_img)

        query_img = query_img.astype('float32')
        query_img /= 255
        reference_img = reference_img.astype('float32')
        reference_img /= 255

        get_feature_layer = K.function([model.layers[0].input] + [K.learning_phase()], [model.layers[-2].output])

        print('inference start')

        # inference
        query_vecs = get_feature_layer([query_img, 0])[0]

        # caching db output, db inference
        db_output = './db_infer.pkl'
        if os.path.exists(db_output):
            with open(db_output, 'rb') as f:
                reference_vecs = pickle.load(f)
        else:
            reference_vecs = get_feature_layer([reference_img, 0])[0]
            with open(db_output, 'wb') as f:
                pickle.dump(reference_vecs, f)

        # l2 normalization
        query_vecs = l2_normalize(query_vecs)
        reference_vecs = l2_normalize(reference_vecs)

        # Calculate cosine similarity
        sim_matrix = np.dot(query_vecs, reference_vecs.T)

        retrieval_results = {}

        for (i, query) in enumerate(queries):
            query = query.split('/')[-1].split('.')[0]
            sim_list = zip(references, sim_matrix[i].tolist())
            sorted_sim_list = sorted(sim_list, key=lambda x: x[1], reverse=True)

            ranked_list = [k.split('/')[-1].split('.')[0] for (k, v) in sorted_sim_list]  # ranked list

            retrieval_results[query] = ranked_list
        print('done')

        return list(zip(range(len(retrieval_results)), retrieval_results.items()))

qfeatures = torch.zeros(len(queries), 512, dtype=torch.float)
endbatch = len(queries) // batch_size
if len(queries) % batch_size != 0:
endbatch += 1
for batidx in range(endbatch):
st = batidx * batch_size
en = min((batidx + 1) * batch_size, len(queries))
xbatch = query_img[st: en, :, :, :]
qfeatures[st: en, :] = feature_model(xbatch.to(device))

Я искал в Google и нашел хороший ответ, но я не пользователь tenforflow ... Я новичок в этой области и использовал только Керас. Я хочу разделить образцы на несколько партий, чтобы поместить их в модель. Как я могу представить разделение кода в Python Keras для исправления infer ()? Я застрял с этим на три дня ... Пожалуйста, помогите мне ... Я действительно ценю за вас заранее ..

...