炼丹琐碎知识点

pin_memory

try_gpu() (来自于李沐老师《动手学深度学习》)

#check device
def get_device():
  return 'cuda' if torch.cuda.is_available() else 'cpu'

def try_gpu(i=0): 
    """如果存在，则返回gpu(i)，否则返回cpu()。"""
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

def try_all_gpus(): 
    """返回所有可用的GPU，如果没有GPU，则返回[cpu(),]。"""
    devices = [torch.device(f'cuda:{i}')
             for i in range(torch.cuda.device_count())]
    return devices if devices else [torch.device('cpu')]

# get device 
#device = get_device()
device = try_gpu(5) # 申请多少号gpu
print(f'DEVICE: {device}')

torch.argmax()

三维的，第一个维度是z，然后x,y同二维！
argmax是返回下标，max是返回 值

【搞懂PyTorch】torch.argmax() 函数详解

torchvision.transforms

官方
PyTorch 学习笔记（三）：transforms的二十二个方法
操作效果图

数据增强的效果展示

transforms.RandomResizedCrop()

保存模型参数torch.save

一文梳理pytorch保存和重载模型参数攻略