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显存分配机制与显存占用分析方法

你可能感兴趣的:(pytorch,深度学习,神经网络)