Keras- Tensorflow Backend에서 특정 디바이스 사용법

Keras- Tensorflow Backend에서 특정 디바이스 사용법

Keras에 tensorflow backend를 사용할 때, 기본적으로 gpu를 사용하여 돌아갑니다. gpu가 있는데 cpu를 사용할 필요는 없겠지만, gpu의 성능을 측정하거나 하는 목적을 위해 강제로 cpu를 사용하도록 하는 방법에 대해 알아보겠습니다.

keras + tensorflow에서 사용가능한 하드웨어 자원보기

아래와 같이 tensorflow의 device_lib이라는 함수를 통해 현재 사용가능한 하드웨어 디바이스들의 리스트를 볼 수 있습니다. 여기서 name을 통해 해당 하드웨어에 접근할 수 있습니다. 제가 사용하는 환경의 경우 1080ti 4개를 사용하는데, 이 디바이스들의 리스트가 나오는 것을 볼 수 있습니다.

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())

[name: "/cpu:0"
device_type: "CPU"
memory_limit: 268435456
locality {
}
incarnation: ....
, name: "/gpu:0"
device_type: "GPU"
memory_limit: 295043072
locality {
  bus_id: 1
}
incarnation: .....
physical_device_desc: "device: 0, name: GeForce GTX 1080 Ti, pci bus id: 0000:02:00.0"
, name: "/gpu:1"
device_type: "GPU"
memory_limit: 366346240
locality {
  bus_id: 1
}
...

특정 device 사용하여 모델 구축, 트레이닝하기

아래와 같이 keras.backend.tensorflow_backend를 K 라는 이름으로 불러오고 with 문을 통해 해당 환경에서 모델을 구축후, fitting을 하면 해당 디바이스 환경에서 모델을 트레이닝 할 수 있습니다. 주의할 점은 model = Sequential()을 통해 모델을 생성할 때부터 with문을 통해 해당 device를 사용한다는 것을 알려줘야 한다는 것입니다. 단순히 model.fit 만 with 아래에서 돌리면 with 문이 먹히지 않더군요..

import keras.backend.tensorflow_backend as K
with K.tf.device('/gpu:0'):
    model = Sequential()
    model.add(Dense(512, input_dim=28*28, activation='relu'))
    model.add(Dense(256, activation='relu'))
    model.add(Dense(128, activation='relu'))
    model.add(Dense(32, activation='relu'))
    model.add(Dense(10, activation='softmax'))

    model.compile(loss='categorical_crossentropy',
                  optimizer='rmsprop',
                  metrics=['accuracy'])
    h = model.fit(X_train, y_train_cat, batch_size=128*4, epochs=10, verbose=1, validation_split=0.3)

Train on 42000 samples, validate on 18000 samples
Epoch 1/10
42000/42000 [==============================] - 0s - loss: 0.5832 - acc: 0.8240 - val_loss: 0.4741 - val_acc: 0.8254
Epoch 2/10
42000/42000 [==============================] - 0s - loss: 0.2057 - acc: 0.9379 - val_loss: 0.1414 - val_acc: 0.9566
Epoch 3/10
42000/42000 [==============================] - 0s - loss: 0.1359 - acc: 0.9585 - val_loss: 0.1454 - val_acc: 0.9548
Epoch 4/10
42000/42000 [==============================] - 0s - loss: 0.0919 - acc: 0.9717 - val_loss: 0.0964 - val_acc: 0.9716
Epoch 5/10

gpu를 사용하면 아래와 같이 한 epoch 당 1초도 안되어서 트레이닝이되는 것을 볼 수 있습니다. 이번엔 cpu로 한 번 모델을 구축해보겠습니다. 아래와 같이 device 이름을 /cpu:0으로 변경한 후 모델을 구축, fitting을 하면, 한 epoch 당 2초가 걸리는 것을 확인할 수 있습니다.

import keras.backend.tensorflow_backend as K
with K.tf.device('/cpu:0'):
    model = Sequential()
    model.add(Dense(512, input_dim=28*28, activation='relu'))
    model.add(Dense(256, activation='relu'))
    model.add(Dense(128, activation='relu'))
    model.add(Dense(32, activation='relu'))
    model.add(Dense(10, activation='softmax'))

    model.compile(loss='categorical_crossentropy',
                  optimizer='rmsprop',
                  metrics=['accuracy'])
    h = model.fit(X_train, y_train_cat, batch_size=128*4, epochs=10, verbose=1, validation_split=0.3)

Train on 42000 samples, validate on 18000 samples
Epoch 1/10
42000/42000 [==============================] - 2s - loss: 0.6009 - acc: 0.8115 - val_loss: 0.3323 - val_acc: 0.8966
Epoch 2/10
42000/42000 [==============================] - 2s - loss: 0.2008 - acc: 0.9404 - val_loss: 0.3604 - val_acc: 0.8832
Epoch 3/10
42000/42000 [==============================] - 2s - loss: 0.1293 - acc: 0.9605 - val_loss: 0.1384 - val_acc: 0.9578
Epoch 4/10
42000/42000 [==============================] - 2s - loss: 0.0914 - acc: 0.9716 - val_loss: 0.1376 - val_acc: 0.9590
Epoch 5/10
42000/42000 [==============================] - 2s - loss: 0.0676 - acc: 0.9796 - val_loss: 0.1145 - val_acc: 0.9674
Epoch 6/10
42000/42000 [==============================] - 2s - loss: 0.0494 - acc: 0.9842 - val_loss: 0.0903 - val_acc: 0.9742