mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-11 12:31:23 +08:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
babaa088d7 | ||
|
|
8217ebe029 | ||
|
|
9311812829 | ||
|
|
2357815640 |
@@ -40,13 +40,32 @@ In this setting, the model will run on all 8 GPUs at once using DataParallel und
|
||||
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
|
||||
|
||||
# DEFAULT
|
||||
|
||||
trainer = Trainer(gpus=[0,1,2,3,4,5,6,7])
|
||||
```
|
||||
|
||||
---
|
||||
#### Multi-node
|
||||
COMING SOON.
|
||||
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
|
||||
|
||||
@@ -150,14 +150,20 @@ class Trainer(TrainerIO):
|
||||
self.use_ddp = False
|
||||
self.use_dp = False
|
||||
|
||||
|
||||
# gpus come in as a string.
|
||||
# if gpus = -1 then use all available devices
|
||||
# otherwise, split the string using commas
|
||||
if gpus is not None:
|
||||
if gpus == '-1':
|
||||
self.data_parallel_device_ids = list(range(0, torch.cuda.device_count()))
|
||||
if type(gpus) is list:
|
||||
self.data_parallel_device_ids = gpus
|
||||
elif type(gpus) is str:
|
||||
if gpus == '-1':
|
||||
self.data_parallel_device_ids = list(range(0, torch.cuda.device_count()))
|
||||
else:
|
||||
self.data_parallel_device_ids = [int(x.strip()) for x in gpus.split(',')]
|
||||
else:
|
||||
self.data_parallel_device_ids = [int(x.strip()) for x in gpus.split(',')]
|
||||
raise Exception('gpus has to be a string or list of ids')
|
||||
|
||||
# set the correct cuda visible devices (using pci order)
|
||||
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
|
||||
@@ -171,6 +177,15 @@ class Trainer(TrainerIO):
|
||||
self.use_dp = distributed_backend == 'dp'
|
||||
self.use_ddp = distributed_backend == 'ddp'
|
||||
|
||||
# use ddp automatically if nb_gpu_nodes > 1
|
||||
if nb_gpu_nodes > 1:
|
||||
self.use_ddp = True
|
||||
self.use_ddp = False
|
||||
w = 'DataParallel does not support nb_gpu_nodes > 1. ' \
|
||||
'Switching to DistributedDataParallel for you. ' \
|
||||
'To silence this warning set distributed_backend=ddp'
|
||||
warnings.warn(w)
|
||||
|
||||
# process info
|
||||
self.proc_rank = 0
|
||||
|
||||
|
||||
@@ -7,7 +7,7 @@ from setuptools import setup, find_packages
|
||||
# http://blog.ionelmc.ro/2014/05/25/python-packaging/
|
||||
setup(
|
||||
name="pytorch-lightning",
|
||||
version='0.2.6',
|
||||
version='0.3.1',
|
||||
description="The Keras for ML researchers using PyTorch",
|
||||
author="William Falcon",
|
||||
author_email="waf2107@columbia.edu",
|
||||
|
||||
Reference in New Issue
Block a user