TMUX: Basic Usage for ML Practitioners
In this article, we explore how you can get started using tmux for machine learning, and never lose your secure shell (ssh) or remote work ever again.
Created on July 7|Last edited on June 12
Comment
During our recent MLOPS course, we showed people how to iterate fast on their machine learning journey. We focused mostly on how to make better use of Weights & Biases in your pipelines for model development, training, and validation.
But we also used a bunch of other tools in our day-to-day work like git, github, jupyter, conda, docker. In lesson 2, I was performing hyperparameter optimization, and to achieve this, we launched a bunch of training runs on a remote machine. In other words: we had to connect, run some code, and orchestrate our experimenting.
An essential tool in this type of work for me is tmux. This handy terminal-based tool enables the user to set up a persistent terminal session on a remote (or local) machine. This is crucial when dealing with remote machines, as the ssh connection is somewhat prone to failing. I actually (accidentally!) closed my terminal session during the course, but as we were using tmux, we could quickly get back on track (thank you, tmux). The tmux session will simply be disconnected when your ssh connection is lost, but all the processes that are active in your session will continue to run in the background on the server. Simply ssh back into the server and attach to the active tmux session to continue work!
Tmux is way more than that, though. It's a handy tool to have even on local machines. You can create multiple panels (Panes) and split views as you see in the image below. Sometimes I would create a session for a specific work, and I know that I can come back and reattach whenever I need to continue this work, knowing all the terminal windows and programs I was running will be there waiting for me.
An excellent resource on tmux and other terminal based tools is the MIT CSAIL Missing Semester course.
Here's what we'll be covering in this article:
Table of Contents
Getting Started
First, you will need to install tmux on the machine you want to use tmux on. In my case, the remote machine that I will be connecting to.
We will use Linux as this is what most people configure to use on remote/cloud machines
💡
# on bash/szh$ sudo apt install tmux
There's also a simple command to check that tmux is working:
$ tmux
Once inside you get a somewhat rich terminal window, that has various informations for you on the screen:
- This is the numbering of the [session: window, pane]. In this case, we are in session 0, in the first window, and the first panel. (we only have one of each, I don't know why session numbering starts from 0 🤣)
- 1:szh* The * is telling you that you are in window 1 named zsh
By default, the tmux prefix command is Ctrl + b (you can change this in your tmux.conf file, I use Ctrl + a). Once the tmux prefix is issued, you can execute a tmux action. For example:
- c will create a new window.
- " will split the current window on 2 vertically (you see Shift + '
- % split vertically (you will see Shift + 5 )
- ⬅️➡️ arrow keys to move across panels
- , rename a window
- exit or Ctrl + d to close a window/panel
Whenever you create a new panel, or split a window. Your cursor will be moved to that window.
I use Ctrl + a as tmux prefix instead of Ctrl + b
You can also create new windows (they look more like tabs anyway). To navigate between windows, you will use p (previous) and n (next). I usually just hit n multiple times and circle back. If you have a lot of windows, you can also use the window number and jump directly to that one.
Reconnecting to tmux
This is why we are using tmux anyway. How do I go back to tmux once my session has been disconnected?
TLDR: If you have only one tmux session running, run $ tmux a
💡
That's it! You got your hands on basic tmux usage.
Running multiple sweep agents on a multi-GPU machine
Whenever you launch computations on a remote machine, you don't want them to be interrupted by ssh connections or internet disconnects. So, tmux to the rescue!
In Lesson 2 of our MLOPs course, I showed how to launch 2 agents in parallel on the same machine. My machine has 2 GPUS, so I can launch two instances of my training script in parallel and perform the sweep twice as fast. Let's do this using tmux:
Launch the sweep controller in the first window
Launch 2 agents in parallel
We will create 2 agents, one per GPU. Also, we monitor nvidia-smi to make sure both GPU are in use.
- We copy the wandb agent command
- We create one extra window with Ctrl + a + c and paste the command
- We do the same thing, but with CUDA_VISIBLE_DEVICES=1 to launch the second agent on the 2nd GPU (GPUs 0 and 1). Default is GPU 0.
- We create one last window to check the nvidia-smi and see if both GPUs are being used.
You have to prefix you wandb agent call with CUDA_VISIBLE_DEVICES=<num GPU>.
💡
I use the alias ca = conda activate in this screen capture
How to Perform Massive Hyperparameter Experiments with W&B
In this article, we look at how Weights & Biases can help you track and organize a massive experiment with thousands of runs.
Leveraging Pre-Trained Models for Image Classification
In this article, we fine-tune a pre-trained model on a new classification dataset, to understand how well transfer learning helps the model train on new data.
Conclusion
That's it for today's brief look at a few things you can do with tmux. If you'd like to check out the course, you can sign up here, and we'll send it to you soon as it's ready. Thanks!
Add a comment
Iterate on AI agents and models faster. Try Weights & Biases today.