mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-20 13:10:42 +08:00
75 lines
2.1 KiB
Markdown
75 lines
2.1 KiB
Markdown
Lightning makes multi-gpu training and 16 bit training trivial.
|
|
|
|
*Note:*
|
|
None of the flags below require changing anything about your lightningModel definition.
|
|
|
|
---
|
|
#### 16-bit mixed precision
|
|
16 bit precision can cut your memory footprint by half. If using volta architecture GPUs it can give a dramatic training speed-up as well.
|
|
First, install apex (if install fails, look [here](https://github.com/NVIDIA/apex)):
|
|
```bash
|
|
$ git clone https://github.com/NVIDIA/apex
|
|
$ cd apex
|
|
$ pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
|
|
```
|
|
|
|
then set this use_amp to True.
|
|
``` {.python}
|
|
# DEFAULT
|
|
trainer = Trainer(amp_level='O2', use_amp=False)
|
|
```
|
|
|
|
---
|
|
#### Single-gpu
|
|
Make sure you're on a GPU machine.
|
|
```python
|
|
# set these flags
|
|
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
|
|
|
|
# DEFAULT
|
|
trainer = Trainer(gpus=[0])
|
|
```
|
|
|
|
---
|
|
#### multi-gpu
|
|
Make sure you're on a GPU machine. You can set as many GPUs as you want.
|
|
In this setting, the model will run on all 8 GPUs at once using DataParallel under the hood.
|
|
```python
|
|
# set these flags
|
|
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
|
|
|
|
|
|
trainer = Trainer(gpus=[0,1,2,3,4,5,6,7])
|
|
```
|
|
|
|
---
|
|
#### Multi-node
|
|
Multi-node training is easily done by specifying these flags.
|
|
```python
|
|
# train on 12*8 GPUs
|
|
trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], nb_gpu_nodes=12)
|
|
```
|
|
|
|
In addition, make sure to set up your SLURM job correctly via the [SlurmClusterObject](https://williamfalcon.github.io/test-tube/hpc/SlurmCluster/). In particular, specify the number of tasks per node correctly.
|
|
|
|
```python
|
|
cluster = SlurmCluster(
|
|
hyperparam_optimizer=test_tube.HyperOptArgumentParser(),
|
|
log_path='/some/path/to/save',
|
|
)
|
|
|
|
# configure cluster
|
|
cluster.per_experiment_nb_nodes = 12
|
|
cluster.per_experiment_nb_gpus = 8
|
|
|
|
cluster.add_slurm_cmd(cmd='ntasks-per-node', value=8, comment='1 task per gpu')
|
|
```
|
|
|
|
---
|
|
#### Self-balancing architecture
|
|
Here lightning distributes parts of your module across available GPUs to optimize for speed and memory.
|
|
|
|
COMING SOON.
|