Use multiple GPUs during predict

[darrylbobo]

I am trying to extract middle layer activation maps using a trained model.
However, when I set to 2 GPUs, it always just use only 1 GPU, and the other GPU is 0%.
.
From this thread #4968, it said I should set kv-store to "device", instead of "local"

However, since I didn't fit the model to new data, all I used is predict function. Where should I change the kv-store ???

Here are the code:

model_load = mx.model.FeedForward.load(prefix, 0, ctx=[mx.gpu(0), mx.gpu(1)],numpy_batch_size=1)

layer_name = ['relu1_2_output','relu2_2_output','relu3_3_output','relu4_3_output', 'relu5_3_output' , 'relu6_output', 'relu7_output']

all_layers = model_load.symbol.get_internals()
fea_symbol = all_layers[layer_name[2]]

feature_extractor = mx.model.FeedForward( ctx=[mx.gpu(0), mx.gpu(1)], symbol=fea_symbol,
numpy_batch_size=1, arg_params=model_load.arg_params,
aux_params=model_load.aux_params,
allow_extra_params=True)

[val_feature, valdata, vallabel]= feature_extractor.predict(img, return_data=True)

[piiswrong]

use module instead

[szha]

This issue is closed due to lack of activity in the last 90 days. Feel free to ping me to reopen if this is still an active issue. Thanks!