pytorch GPU训练

1、查看cuda相关

    1）cuda是否可用    

        torch.cuda.is_available()
    2）gpu数量

        torch.cuda.device_count()

    3）gpu名字，默认从0开始

        torch.cuda.get_device_name(0)
    4）当前设备索引

        torch.cuda.current_device()

2、多卡训练

    1）DataParallel

    2）

3、设置默认device

    1）torch.cuda.set_device(device_ids)

    2）os.environ["CUDA_VISIBLE_DEVICES"] = device_ids

【注意】

        1）当已经import torch并进行了初始化之后，2可能无效，导致模型在其他卡上训练

        2）torch.load(checkpoint)，如果checkpoint中的tensor是gpu tensor，则会加载到0卡上；

            对于1方法来说，是绝对0卡gpu: 0；

            对于2方法来说，是相对0卡gpu: visible_start；

            最好是存储checkpoint时，存为cpu tensor；或者load时，设置map_location=cpu

        3）不推荐使用model.to(device)来指定模型gpu，因为不设置全局的默认设备的话，torch.cuda.empty_cache()操作可能会出现问题

4、GPU显存限制

[Feature] Allow user to specify a fraction of the GPU memory. (#48172) · pytorch/pytorch@47aa253 · GitHub

torch.cuda.set_per_process_memory_fraction — PyTorch 1.12 documentation

Pytorch显存分配机制与显存占用分析方法