使用gpus设备训练
这个也很简单,但要注意版本,官方文档给出的是2.0以上的使用方法。使用方法是设置Trainer的accelerator和device 参数。
# 使用k个gpu
# DEFAULT (int) specifies how many GPUs to use per node
Trainer(accelerator="gpu", devices=k)
# Above is equivalent to
Trainer(accelerator="gpu", devices=list(range(k)))
# 使用指定的两块gpu
# Specify which GPUs to use (don't use when running on cluster)
Trainer(accelerator="gpu", devices=[0, 1])
# 和上面一致
# Equivalent using a string
Trainer(accelerator="gpu", devices="0, 1")
# 使用所有的gpu
# To use all available GPUs put -1 or '-1'
# equivalent to list(range(torch.cuda.device_count()))
Trainer(accelerator="gpu", devices=-1)
训练精度设置
当显存、速度有额外要求时,可以牺牲部分精度换取最少的存储消耗和更快的速度。指定Trainer的precision参数。设置后lightning自动将tensor加载至相应设备,即不用调用.to(device)方法。
# 16位精度
Trainer(precision='16-mixed')
# 32位精度
Trainer(precision="32-true")
# or
Trainer(precision="32")
# or
Trainer(precision=32)
# 64位精度
Trainer(precision="64-true")
# or
Trainer(precision="64")
# or
Trainer(precision=64)
Precision | CPU | GPU | TPU | IPU |
---|---|---|---|---|
16 Mixed | No | Yes | No | Yes |
BFloat16 Mixed | Yes | Yes | Yes | No |
32 True | Yes | Yes | Yes | Yes |
64 True | Yes | Yes | No | No |