# argument
parser.add_argument('--ngpu', type=int, default=1, help='0 = CPU.')
parser.add_argument('--gpu_id', type=int, default=0,
help='device range [0,ngpu-1]')
parser.add_argument('--workers', type=int, default=4,
help='number of data loading workers (default: 2)')
--ngpu
: ngpu嘛n为数量,默认为1。通过这个参数,可以选择在GPU上进行计算还是使用CPU进行计算。当–ngpu为0时,表示不使用GPU,即使用CPU进行计算。gpu:0是主cpu
--gpu_id
:指定要使用的GPU设备的ID,默认为0。区间为:[0, ngpu - 1]
--workers
:工作线程(workers)是指在数据加载过程中同时运行的线程数量。每个工作线程负责从数据集中加载一部分数据,并将其放入内存中供模型使用
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
if args.ngpu == 1:
# make only device #gpu_id visible, then
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_id)
这段代码首先设置了环境变量CUDA_DEVICE_ORDER
的值为"PCI_BUS_ID"。这是为了告诉系统按照PCI总线ID的顺序对GPU设备进行排序,确保GPU设备的顺序与物理插槽上的顺序一致。然后调用os.environ
模块,用于访问和修改操作系统的环境变量。这在某些情况下可能很有用,例如当系统中有多个GPU设备,并且希望将训练过程限制在某一个特定的GPU上,或者在单GPU环境下确保不会占用其他GPU资源。
设置GPU设备对程序的可见性:在多GPU环境下,可以通过设置CUDA_VISIBLE_DEVICES
环境变量来控制哪些GPU设备对程序可见。
设置设备可见性的原理就是设置环境变量:除了通过os.environ
,你可以获取当前操作系统的所有环境变量,并可以根据需要进行读取和修改。它返回一个字典对象,其中键是环境变量的名称,值是对应的值之外,也可以写bash文件.sh,如下:
export CUDA_VISIBLE_DEVICES=“0,3”
args.use_cuda = args.ngpu > 0 and torch.cuda.is_available() # check GPU
if args.use_cuda:
torch.cuda.manual_seed_all(args.manualSeed)
这意味着每个GPU上的随机数生成器都使用相同的种子,从而使得在并行训练中每个GPU的计算顺序和结果保持一
致,从而实现结果的可重现性
cudnn.benchmark = True
通过将cudnn.benchmark设置为True,可以启用自动寻找最适合当前配置的卷积算法,从而提高运行速度
if args.use_cuda:
if args.ngpu > 1:
net = torch.nn.DataParallel(net, device_ids=list(range(args.ngpu)))
nn.DataParallel
会自动管理多个GPU设备的可见性。
nn.DataParallel对模型进行包装,使其能够同时在多个GPU上运行。nn.DataParallel会自动将输入数据划分为多个子批次,并将每个子批次分配到不同的GPU上进行处理。然后,它会自动将每个GPU上计算的结果进行收集、合并和同步,最终返回整个批次的输出结果
来自torch.optim:具有通用优化算法(如SGD,Adam等)的优化包 & torch.optim.lr_scheduler接口实现
optimizer.zero_grad()
因为训练的过程通常使用mini-batch方法,所以如果不将梯度清零的话,梯度会与上一个batch的数据相关,因此该函数要写在反向传播和梯度下降之前。loss.backward()
:PyTorch的反向传播(即tensor.backward())是通过autograd包来实现的,autograd包会根据tensor进行过的数学运算来自动计算其对应的梯度。如果你设置tensor的requires_grads为True,就会开始跟踪这个tensor上面的所有运算,如果你做完运算后****使用tensor.backward(),所有的梯度就会自动运算,tensor的梯度将会累加到它的.grad属性里面去。step()
函数的作用是执行一次优化步骤,通过梯度下降法来更新参数的值。因为梯度下降是基于梯度的,所以在执行optimizer.step()函数前应先执行loss.backward()函数来计算梯度。直接看这个学习用法
def lr_scheduler(optimizer, epoch):
"""Decay learning rate by a factor of 0.1 every lr_decay_epoch epochs."""
lr_list = [100, 140, 240]
if epoch in lr_list:
print('change the learning rate')
for param_group in optimizer.param_groups:
param_group['lr'] = param_group['lr'] * 0.1
return optimizer
criterion = nn.CrossEntropyLoss()
# optimizer = torch.optim.Adam(ann.parameters(), lr=learning_rate)
optimizer = torch.optim.SGD(ann.parameters(), lr=learning_rate, momentum=0.9, weight_decay=5e-4)
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(tqdm(train_loader), 0):
# 将 trainloader 包装成一个可迭代对象,并在每次迭代时显示进度条。在循环中使用 tqdm(trainloader) 可以方便地看到当前迭代的进度
# 第二个参数是可选的 start 参数,用于指定索引的起始值,默认为 0
ann.train()
ann.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
optimizer.zero_grad() # 梯度归零
loss.backward() # 计算梯度
optimizer.step() # 更新参数
optimizer = lr_scheduler(optimizer, epoch) #每个epoch更新一次lr参数