[Python] 딥러닝 학습하는 도중에 GPU 사용량 확인하기

이번 글은 실제 학습 동안에, GPU를 얼마나 사용하는지를 프린트해주는 것을 찾게 되어서 공유한다.

실제로는 nvidia-smi를 통해서 확인할 수 있지만, 도커를 쓰다보면 이러한 사용량을 알 수 없는 경우가 생긴다.

그래서 이번에는 학습하면서 중간 중간에 출력을 해줘서 실제로 사용 여부를 확인할 수가 있다.

 

추가적으로 threading을 사용해서 학습하면서, 병렬로 돌릴 수 있는 코드를 발견해서 공유한다.

 

# Import MNIST data
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("./", one_hot=False)
import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np

# Parameters
learning_rate = 0.1
num_steps = 10000
batch_size = 128
display_step = 100

# Network Parameters
n_hidden_1 = 256 # 1st layer number of neurons
n_hidden_2 = 256 # 2nd layer number of neurons
num_input = 784 # MNIST data input (img shape: 28*28)
num_classes = 10 # MNIST total classes (0-9 digits)

아래 코드가 GPUtil에 있는 코드이다.

해당 코드를 돌리면 주기적으로 몇 초 간격으로 사용량을 보여준다.

import GPUtil
from threading import Thread
import time

class Monitor(Thread):
    def __init__(self, delay):
        super(Monitor, self).__init__()
        self.stopped = False
        self.delay = delay # Time between calls to GPUtil
        self.start()

    def run(self):
        while not self.stopped:
            GPUtil.showUtilization()
            time.sleep(self.delay)

    def stop(self):
        self.stopped = True
        
monitor = Monitor(10)
#monitor.stop()

이런 식으로 프린트가 된다. 

텐서플로우 모델 아키텍처를 구성한다. 

input_fn = tf.estimator.inputs.numpy_input_fn(
    x={'images': mnist.train.images}, y=mnist.train.labels,
    batch_size=batch_size, num_epochs=None, shuffle=True)

# Define the neural network
def neural_net(x_dict):
    x = x_dict['images']
    layer_1 = tf.layers.dense(x, n_hidden_1)
    layer_2 = tf.layers.dense(layer_1, n_hidden_2)
    out_layer = tf.layers.dense(layer_2, num_classes)
    return out_layer
def model_fn(features, labels, mode):
    logits = neural_net(features)
    pred_classes = tf.argmax(logits, axis=1)
    pred_probas = tf.nn.softmax(logits)
    loss_op = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
        logits=logits, labels=tf.cast(labels, dtype=tf.int32)))
    optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate)
    train_op = optimizer.minimize(loss_op, global_step=tf.train.get_global_step())
    acc_op = tf.metrics.accuracy(labels=labels, predictions=pred_classes)
    estim_specs = tf.estimator.EstimatorSpec(
      mode=mode,
      predictions=pred_classes,
      loss=loss_op,
      train_op=train_op,
      eval_metric_ops={'accuracy': acc_op})

    return estim_specs
model = tf.estimator.Estimator(model_fn)

실제로 학습시키면서 계속 프린트가 된다.

model.train(input_fn, steps=num_steps)

멈추려면 아래 코드를 실행하면 된다.

monitor.stop()

위의 예제처럼 threaing 라이브러리와 GPUtil 라이브러리를 사용해서 학습을 하면서도, 동시에 GPU 사용량을 주기적으로 프린트할 수 있다.