这篇教程Pytorch dataloader在加载最后一个batch时卡死的解决写得很实用,希望能帮到您。 问题:自己写了个dataloader,为了部署方便,用OpenCV的接口进行数据读取,而没有用PIL,代码大致如下: def __getitem__(self, idx): sample = self.samples[idx] img = cv2.imread(sample[0]) img = cv2.resize(img, tuple(self.input_size)) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # if not self.val and random.randint(1, 10) < 3: # img = self.img_aug(img) img = Image.fromarray(img) img = self.transforms(img) ... 结果在训练过程中,在第1个epoch的最后一个batch时,程序卡死。 解决方案:可能是因为OpenCV与Pytorch互锁的问题,关闭OpenCV的多线程,问题解决。 cv2.setNumThreads(0)cv2.ocl.setUseOpenCL(False) 补充:pytorch 中一个batch的训练过程 # 一般情况下optimizer.zero_grad() # 梯度清零preds = model(inputs) # inference,前向传播求出预测值loss = criterion(preds, targets) # 计算lossloss.backward() # 反向传播求解梯度optimizer.step() # 更新权重,更新网络权重参数 此外,反向传播前,如果不进行梯度清零,则可以实现梯度累加,从而一定程度上解决显存受限的问题。 以上为个人经验,希望能给大家一个参考,也希望大家多多支持51zixue.net。 Python趣味挑战之用pygame实现飞机塔防游戏 Python基础学习之条件控制语句小结 |