7.pytorch lightning 之GPU设置

使用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