GPU underutilized in Actor Critic (A2C) Stable Baselines3 implementation

user_3748104 0
Viewed 2101

I am trying to use A2C of StablesBaselines3 for training an agent on my custom environment. My problem is that my GPU Utilization is very less (around 10 % only) while my CPU utilization has hit the ceiling. Because of this the training is very very slow. I have tried the following things as per this discussion thread addressed by @araffin [https://github.com/hill-a/stable-baselines/issues/308]

  • Expanding the Actor and Critic architecture to a three layer neural network having 256, 256 and 128 neurons respectively. The GPU utilization did increase after that but it was only marginal (increased from 10 % to 15 %) as in this suggestion
  • changed device argument of A2C method to 'cuda' from the default which is 'auto' - No improvement

It is not entirely clear what worked for the user in the above mentioned discussion. My A2C is intialized as follows:

import gym
from stable_baselines3 import A2C
policy_kwargs = dict(activation_fn=th.nn.ReLU,
                     net_arch=[dict(pi=[256, 256, 128], vf=[256, 256, 128])])
model = A2C("MlpPolicy", env, verbose=1,  learning_rate=linear_schedule(0.001),\
 n_steps=50, gamma=0.8, gae_lambda=1.0, ent_coef=0.2, vf_coef=0.2, \
tensorboard_log=logdir,  policy_kwargs=policy_kwargs, seed=50, device='cuda' )
model.learn(total_timesteps=50000)

GPU configurations are:

  • Cuda version - 9.0
  • GPU - NVIDIA RTX 2080i
  • GPU mem - 11GB (out of that only about 1600MB is being used)

CPU configurations

  • Threads per Core - 2
  • Core per socket - 8
  • Intel i9 core 9900 K @ 3.6 Ghz
  • CPU Mem - 32GB SW versions
  • Torch - 1.6.0
  • Tensorflow GPU - 2.4.1 (although I know SB3 implementation uses Torch)
  • Python 3.7 on anaconda environment

Any help will be appreciated!

reinforcement-learning pytorch gpu
1 Answers

Stable baselines is using your gpu ... if you look task manager on second tab click on your gpu and instead of 3D select cuda and you will see the usage of cuda

I had some troubles with my env because my env uses pandas .. and pandas use cpu .. on windows is not possible easly to use cudf so my cpu was using 100%

I used pandas and numpy on my custom env .. to improve it I should use cunumery and cudf instead, but to do it I need linux.

so for now, my cpu is using 100% gpu is using cuda 70% and gpu memory almost 1.7 gb

user_181715170
Related