3.2 基本配置#

对于一个PyTorch项目，我们需要导入一些Python常用的包来帮助我们快速实现功能。常见的包有os、numpy等，此外还需要调用PyTorch自身一些模块便于灵活使用，比如torch、torch.nn、torch.utils.data.Dataset、torch.utils.data.DataLoader、torch.optimizer等等。

经过本节的学习，你将收获：

  • 在深度学习/机器学习中常用到的包

  • GPU的配置

首先导入必须的包。注意这里只是建议导入的包导入的方式，可以采用不同的方案，比如涉及到表格信息的读入很可能用到pandas，对于不同的项目可能还需要导入一些更上层的包如cv2等。如果涉及可视化还会用到matplotlib、seaborn等。涉及到下游分析和指标计算也常用到sklearn。

import os 
import numpy as np 
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torch.optim as optimizer

根据前面我们对深度学习任务的梳理，有如下几个超参数可以统一设置，方便后续调试时修改：

  • batch size

  • 初始学习率（初始）

  • 训练次数（max_epochs）

  • GPU配置

batch_size = 16
# 批次的大小
lr = 1e-4
# 优化器的学习率
max_epochs = 100

除了直接将超参数设置在训练的代码里，我们也可以使用yaml、json，dict等文件来存储超参数，这样可以方便后续的调试和修改，这种方式也是常见的深度学习库（mmdetection，Paddledetection，detectron2）和一些AI Lab里面比较常见的一种参数设置方式。

我们的数据和模型如果没有经过显式指明设备，默认会存储在CPU上，为了加速模型的训练，我们需要显式调用GPU，一般情况下GPU的设置有两种常见的方式：

# 方案一：使用os.environ，这种情况如果使用GPU不需要设置
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1' # 指明调用的GPU为0,1号

# 方案二：使用“device”，后续对要使用GPU的变量用.to(device)即可
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu") # 指明调用的GPU为1号

当然还会有一些其他模块或用户自定义模块会用到的参数，有需要也可以在一开始进行设置。