slow DQN with Keras

Viewed 101

I've tried to write a first version of a simple DQN to play the game snake that I've written. So far it's running, but the speed of it puzzles me. It runs on the GPU, but nvidia-smi shows around 2-4% usage. Every training step with replay takes a few seconds. The game itself has a 10x10 field, one reward block and the snake. That gets into the network as a 10x10 numpy array. Below is my DQN class:

import numpy as np
import keras
from collections import deque
import random


class DQN:
    def __init__(self, sess, input_size, lr=0.01, gamma=0.95, mem_size=1000,
                 epsilon_decay=0.999, eps_min=0.005, epsilon=1.0):
        self.activation = keras.activations.relu
        self.optimizer = keras.optimizers.Adadelta
        self.learning_rate = lr

        self.memory = deque(maxlen=mem_size)

        self.epsilon = epsilon
        self.eps_min = eps_min
        self.epsilon_decay = epsilon_decay

        self.sess = sess
        self.gamma = gamma
        self.model = self.cnn(input_size)
        self.target_model = self.cnn(input_size)

    def cnn(self, input_size):
        initializer = keras.initializers.variance_scaling()
        model = keras.Sequential()
        model.add(keras.layers.Conv2D(16, (2, 2), activation=self.activation, input_shape=(input_size, input_size, 1)))
        model.add(keras.layers.Conv2D(32, (3, 3), activation=self.activation))
        model.add(keras.layers.MaxPooling2D())
        model.add(keras.layers.Flatten())
        model.add(keras.layers.Dense(16, kernel_initializer=initializer, activation='relu'))
        model.add(keras.layers.Dense(4, kernel_initializer=initializer))
        model.compile(optimizer=keras.optimizers.Adadelta(learning_rate=self.learning_rate), loss='mean_squared_error')
        return model

    def remember(self, state, action, reward, next_state, done):
        self.memory.append([state, action, reward, next_state, done])

    def replay(self):
        batch_size = 16
        if len(self.memory) < batch_size:
            return

        samples = random.sample(self.memory, batch_size)
        for sample in samples:
            state, action, reward, next_state, done = sample
            target = self.target_model.predict(state, batch_size=len(state))
            if done:
                target[0][action] = reward
            else:
                q_future = max(self.target_model.predict(next_state, batch_size=len(state))[0])
                target[0][action] = reward + q_future * self.gamma
            self.model.fit(state, target, epochs=1, verbose=0, batch_size=16)

    def target_train(self):
        weights = self.model.get_weights()
        target_weights = self.target_model.get_weights()
        for i in range(len(target_weights)):
            target_weights[i] = weights[i]
        self.target_model.set_weights(target_weights)

    def act(self, state):
        self.epsilon *= self.epsilon_decay
        self.epsilon = max(self.epsilon, self.eps_min)
        if np.random.random() < self.epsilon:
            return self.sess.get_random_action()
        return np.argmax(self.model.predict(state)[0])

In my tests, the replay function would need almost all of the time used per loop. Right now, it runs model.predict and model.fit with every state on its own. I read on some website a comment that stated it used vectorized inputs to speed this up.

  • How could I use vectorized input? I found no example of what this is in detail. Can I use it only with fit or also with predict?
  • How could I find the speed bottleneck in my code? It obviously is not using CPU or GPU very effective.
0 Answers
Related