Pytorch dataloader在加载最后一个batch时卡死的解决

2024-04-02 19:04:59 815人浏览独家记忆

Python 官方文档：入门教程 => 点击学习

摘要

问题：自己写了个dataloader，为了部署方便，用OpenCV的接口进行数据读取，而没有用PIL，代码大致如下： def __getitem__(self, idx

问题：

自己写了个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)  # 计算loss
loss.backward()                   # 反向传播求解梯度
optimizer.step()                  # 更新权重，更新网络权重参数

此外，反向传播前，如果不进行梯度清零，则可以实现梯度累加，从而一定程度上解决显存受限的问题。

以上为个人经验，希望能给大家一个参考，也希望大家多多支持编程网。

您可能感兴趣的文档:

--结束END--

本文标题: Pytorch dataloader在加载最后一个batch时卡死的解决

本文链接: https://lsjlt.com/news/127042.html(转载时请注明来源链接)

有问题或投稿请发送至: 邮箱/279061341@qq.com QQ/279061341