PyTorch常用代码段合集

文章预览

来源 | 极市平台 作者 | Jack Stark@知乎 PyTorch最好的资料是官方文档。本文是PyTorch常用代码段，在参考资料[1](张皓：PyTorch Cookbook)的基础上做了一些修补，方便使用时查阅。 1 『基本配置』 导入包和版本查询 import torch import torch.nn as nn import torchvision print(torch.__version__) print(torch.version.cuda) print(torch.backends.cudnn.version()) print(torch.cuda.get_device_name(0)) 可复现性 在硬件设备（CPU、GPU）不同时，完全的可复现性无法保证，即使随机种子相同。但是，在同一个设备上，应该保证可复现性。具体做法是，在程序开始的时候固定torch的随机种子，同时也把numpy的随机种子固定。 np.random.seed(0) torch.manual_seed(0) torch.cuda.manual_seed_all(0) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False 显卡设置 如果只需要一张显卡 # Device configuration device = torch.device( 'cuda' if torch ………………………………