Как заставить тензорный поток использовать все доступные графические процессоры? - PullRequest
Новая
       31

Как заставить тензорный поток использовать все доступные графические процессоры?

0 голосов
/ 26 апреля 2018

У меня кластер из 8 графических процессоров, и когда я запускаю фрагмент кода Tensorflow (вставлен ниже), он использует только один графический процессор вместо всех 8. Я подтвердил это с помощью nvidia-smi.

# Set some parameters
IMG_WIDTH = 256
IMG_HEIGHT = 256
IMG_CHANNELS = 3
TRAIN_IM = './train_im/'
TRAIN_MASK = './train_mask/'
TEST_PATH = './test/'

warnings.filterwarnings('ignore', category=UserWarning, module='skimage')
num_training = len(os.listdir(TRAIN_IM))
num_test = len(os.listdir(TEST_PATH))
# Get and resize train images
X_train = np.zeros((num_training, IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS), dtype=np.uint8)
Y_train = np.zeros((num_training, IMG_HEIGHT, IMG_WIDTH, 1), dtype=np.bool)
print('Getting and resizing train images and masks ... ')
sys.stdout.flush()

#load training images
for count, filename in tqdm(enumerate(os.listdir(TRAIN_IM)), total=num_training):
    img = imread(os.path.join(TRAIN_IM, filename))[:,:,:IMG_CHANNELS]
    img = resize(img, (IMG_HEIGHT, IMG_WIDTH), mode='constant', preserve_range=True)
    X_train[count] = img
    name, ext = os.path.splitext(filename)
    mask_name = name + '_mask' + ext
    mask = cv2.imread(os.path.join(TRAIN_MASK, mask_name))[:,:,:1]
    mask = resize(mask, (IMG_HEIGHT, IMG_WIDTH))
    Y_train[count] = mask

# Check if training data looks all right
ix = random.randint(0, num_training-1)
print(ix)
imshow(X_train[ix])
plt.show()
imshow(np.squeeze(Y_train[ix]))
plt.show()
# Define IoU metric
def mean_iou(y_true, y_pred):
    prec = []
    for t in np.arange(0.5, 1.0, 0.05):
        y_pred_ = tf.to_int32(y_pred > t)
        score, up_opt = tf.metrics.mean_iou(y_true, y_pred_, 2)
        K.get_session().run(tf.local_variables_initializer())
        with tf.control_dependencies([up_opt]):
            score = tf.identity(score)
        prec.append(score)
    return K.mean(K.stack(prec), axis=0)

# Build U-Net model
inputs = Input((IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS))
s = Lambda(lambda x: x / 255) (inputs)
width = 64
c1 = Conv2D(width, (3, 3), activation='relu', padding='same') (s)
c1 = Conv2D(width, (3, 3), activation='relu', padding='same') (c1)
p1 = MaxPooling2D((2, 2)) (c1)

c2 = Conv2D(width*2, (3, 3), activation='relu', padding='same') (p1)
c2 = Conv2D(width*2, (3, 3), activation='relu', padding='same') (c2)
p2 = MaxPooling2D((2, 2)) (c2)

c3 = Conv2D(width*4, (3, 3), activation='relu', padding='same') (p2)
c3 = Conv2D(width*4, (3, 3), activation='relu', padding='same') (c3)
p3 = MaxPooling2D((2, 2)) (c3)

c4 = Conv2D(width*8, (3, 3), activation='relu', padding='same') (p3)
c4 = Conv2D(width*8, (3, 3), activation='relu', padding='same') (c4)
p4 = MaxPooling2D(pool_size=(2, 2)) (c4)

c5 = Conv2D(width*16, (3, 3), activation='relu', padding='same') (p4)
c5 = Conv2D(width*16, (3, 3), activation='relu', padding='same') (c5)

u6 = Conv2DTranspose(width*8, (2, 2), strides=(2, 2), padding='same') (c5)
u6 = concatenate([u6, c4])
c6 = Conv2D(width*8, (3, 3), activation='relu', padding='same') (u6)
c6 = Conv2D(width*8, (3, 3), activation='relu', padding='same') (c6)

u7 = Conv2DTranspose(width*4, (2, 2), strides=(2, 2), padding='same') (c6)
u7 = concatenate([u7, c3])
c7 = Conv2D(width*4, (3, 3), activation='relu', padding='same') (u7)
c7 = Conv2D(width*4, (3, 3), activation='relu', padding='same') (c7)

u8 = Conv2DTranspose(width*2, (2, 2), strides=(2, 2), padding='same') (c7)
u8 = concatenate([u8, c2])
c8 = Conv2D(width*2, (3, 3), activation='relu', padding='same') (u8)
c8 = Conv2D(width*2, (3, 3), activation='relu', padding='same') (c8)

u9 = Conv2DTranspose(width, (2, 2), strides=(2, 2), padding='same') (c8)
u9 = concatenate([u9, c1], axis=3)
c9 = Conv2D(width, (3, 3), activation='relu', padding='same') (u9)
c9 = Conv2D(width, (3, 3), activation='relu', padding='same') (c9)

outputs = Conv2D(1, (1, 1), activation='sigmoid') (c9)

model = Model(inputs=[inputs], outputs=[outputs])

sgd = optimizers.SGD(lr=0.03, decay=1e-6, momentum=0.9, nesterov=True)
model.compile(optimizer=sgd, loss='binary_crossentropy', metrics=[mean_iou])
model.summary()

# Fit model
earlystopper = EarlyStopping(patience=20, verbose=1)
checkpointer = ModelCheckpoint('nuclei_only.h5', verbose=1, save_best_only=True)
results = model.fit(X_train, Y_train, validation_split=0.05, batch_size = 32, verbose=1, epochs=100, 
                callbacks=[earlystopper, checkpointer])

Я хотел бы использовать mxnet или какой-либо другой метод для запуска этого кода на всех доступных графических процессорах.Однако я не уверен, как это сделать.Все ресурсы показывают, как это сделать на наборе данных mnist.У меня есть собственный набор данных, который я читаю по-другому.Следовательно, не совсем уверен, как изменить код.

1 Ответ

0 голосов
/ 26 апреля 2018

TL; DR : используйте multi_gpu_model() от Keras.

TS; WM :

Из Tensorflow Guide :

Если в вашей системе более одного графического процессора,по умолчанию будет выбран графический процессор с наименьшим идентификатором.

Если вы хотите использовать несколько графических процессоров, к сожалению, вы должны вручную указать, какие тензоры для каждого графического процессора ставить, например 

with tf.device('/device:GPU:2'):

Дополнительная информация в Tensorflow Guide Использование нескольких графических процессоров .

С точки зрения распределения вашей сети по нескольким графическим процессорам, есть два основных способа сделать это.

  1. Вы распределяете сеть по уровням графических процессоров.Это проще в реализации, но не принесет большого выигрыша в производительности, поскольку графические процессоры будут ждать завершения операции друг друга.

  2. Вы создаете отдельные копии своей сети, называемые «вышками»на каждом графическом процессоре.Когда вы загружаете сеть octuple, вы разбиваете входной пакет на 8 частей и распределяете их.Пусть сеть распространяется вперед, затем суммирует градиенты и выполняет обратное распространение.Это приведет к почти линейному ускорению с количеством графических процессоров.Однако реализовать его гораздо сложнее, поскольку вам также приходится сталкиваться со сложностями, связанными с нормализацией партии, и очень желательно убедиться, что вы правильно рандомизировали свою партию. хороший учебник здесь .Вам также следует ознакомиться с указанным там кодом Inception V3 , чтобы узнать, как его структурировать.Особенно _tower_loss(), _average_gradients() и часть train(), начинающаяся с for i in range(FLAGS.num_gpus):.

Если вы хотите попробовать Keras, он теперь упростил multi-gpuзначительно тренируясь с multi_gpu_model().Он может сделать всю тяжелую работу за вас.

...