How to optimise your Starcraft machine learning model for Google Colaboratory.

5 min readApr 19, 2018

Leverage free GPU compute by taking your Starcraft Machine Learning model into the cloud.

TL;DR: We learn how to optimise a PySC2 machine learning model for Google Colabs. The Colabs notebook we refer to throughout this tutorial can be found here:
https://colab.research.google.com/drive/18asH_hmCjgUUCb-_-SPgPWkK0MYLPAmk

In the last post, we described how we successfully managed to get PySc2 working on Colaboratory,

In this post, we will teach you how to optimise your machine learning model to train & run on Colabs.

We will be using William Xu’s Pytorch implementation of DQN in this example, however the basic principles here can be adapted to any model / agent combination out there, so lets get started.

You can get a (Colabified) version of Will’s code using the following command in your console

Git clone https://github.com/PaulConyngham/DeepRLBootCampLabs.git

First & foremost, by default — your Colabs instance’s GPU is turned off. In order to turn it on you need to go to the menu & click “runtime”

Then click “change runtime” type:

Click the drop down arrow where it says “hardware accelerator”

And change this to GPU,

Your Internet Connection

If you disconnect from Colabs at any point during the training phase of your algorithm, the Virtual Machine (VM) on Google’s side will be reset & all of the work will be lost — so make sure you have a solid internet connection!

Also, as an aside, ensure that your laptop does not “time out” during a screen saver or lock screen, as this can cause Colabs to disconnect too. Run some other application that stops your computer from doing this.

Getting your files onto Colabs.

The way we achieved this was via creating a (Colabified) git hub repository of the original code and simply Git clone ‘ed it onto Colabs. However this is quite slow with regards to pushing updates to the code.

Git clone https://github.com/PaulConyngham/DeepRLBootCampLabs.git

An alternative approach is to create a folder on your Google drive & sync your code there-updating your code on the fly. The files in your Google drive folder can then be downloaded onto Colabs. This is achieved using PyDrive , example code below to grab a file called “dogscats.zip”:

from pydrive.auth import GoogleAuth
from pydrive.drive import GoogleDrive
from google.colab import auth
from oauth2client.client import GoogleCredentials# 1. Authenticate and create the PyDrive client.
auth.authenticate_user()
gauth = GoogleAuth()
gauth.credentials = GoogleCredentials.get_application_default()
drive = GoogleDrive(gauth)dogscats=drive.CreateFile({"id":"19RVUGE8zvuGz2xKnmOFqNR2u_1IffkF65"})#create the file "dogscats.zip" on your local Colabs instance...
dogscats.GetContentFile("dogscats.zip")

RAM considerations for your model.

Next up is RAM considerations, if your model is going to be loading in many training examples or building up massive amounts of replay memory (as in DQN) — bare in mind that your Colabs instance has a limit of 12.4 GB of RAM.

After successfully getting PySc2 to run on Colabs, we ran into a bug where the headless SC2 client kept crashing as we found out in time, that our DQN model was exceeding the RAM limits of Colabs- so keep this in mind when tuning your models too.

Your machine learning model may have a different architecture! but as an example, In our code, we limit our amount of replay memory with the following line of code in base_rl_agent.py, in the BaseRLAgent class:

self._memory = ReplayMemory(50000)

Colabs Run Time & Saving your model’s progress

Google kindly gives you access to a powerful Virtual Machine for 12 hours before it is automatically reset. So in order to save your progress whilst training your model, you need to ensure that , for example, if you are using Deep Learning as your machine learning of choice — when the “training” of the model is complete — it is able to save itself so that your progress can be downloaded.

Feel free to implement this anyway you want, but the way that we do this is by saving the “weights” of our neural network, once a termination signal is sent to the model.

We do this every 11 hours, — you can see this implemented in the “Run” section in our Colabs notebook itself:

!PYTHONPATH=. timeout -s 2 11h python ~/pytorch/sc2_agents/BaseTrainer.py --map=MoveToBeacon --train=True

In the command above the “timeout -s 2 11h” means that a user break command will be sent to the running process (in this case the model) once the training has run for 11 hours.

The Pytorch code below is then used in the model to save the Neural networks weights to a file, if you are using Tensorflow or some other model architecture, find a similar command that does the same:

torch.save(self._Q.state_dict(), self._Q_weights_path)

The neural networks weights are saved to a file called “SC2QAgent”

We then download the neural networks weights from Colabs running the following command in the Colabs notebook itself:

from google.colab import filesfiles.download('./pytorch/data/SC2QAgent')

As an aside, what you could do to continue training is download the weights every 11 hours, and reupload them, load the weights into your model & keep training… :)

Colabs File Structure

Keep in mind when downloading your code to Colabs, that your will have to “Colabify” file references within your code in order to get your model to train.

If you have written any custom classes & are importing them into Colabs in order to run your model you need to reference these custom classes correctly (in your Colabs version of your code )

In our example Basetrainer.py is importing a class called BaseRLAgent:

from pytorch.sc2_agents.base_rl_agent import BaseRLAgent as Agent

This is because BaseRLAgent is stored in the file base_rl_agent.py found in the folder pytorch/sc2_agents/ when downloaded to our Colabs instance.

The dots in the above import command just represent the folder structure to the BaseRLAgent class. You will need to update class references directories if you import any custom classes too.

Rendering

If you would like your model to train on Colabs, you have to set the Render = False flag, as your Virtual machine on Google’s side does not have a screen.

We do this in Basetrainer.py in this line of code:

flags.DEFINE_bool(“render”, False, “Whether to render with pygame.”)

Obviously, if you would like to see the visualised results on your computer at home, install the SC2 client, either from Blizzard or get the headless SC2 from Deepmind here, set render = true and you should be set,

And that’s it!

StarAi is a small team of Developers, Machine Learning Engineers & Researchers based in Sydney, Australia