脚本之家,脚本语言编程技术及教程分享平台!
分类导航

Python|VBS|Ruby|Lua|perl|VBA|Golang|PowerShell|Erlang|autoit|Dos|bat|

服务器之家 - 脚本之家 - Python - Pytorch dataloader在加载最后一个batch时卡死的解决

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

2021-11-16 10:57Wilber529 Python

这篇文章主要介绍了Pytorch dataloader在加载最后一个batch时卡死的解决方案,具有很好的参考价值,希望对大家有所帮助。如有错误或未考虑完全的地方,望不吝赐教

问题:

自己写了个dataloader,为了部署方便,用OpenCV的接口进行数据读取,而没有用PIL,代码大致如下:

?
1
2
3
4
5
6
7
8
9
10
11
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的多线程,问题解决。

?
1
2
cv2.setNumThreads(0)
cv2.ocl.setUseOpenCL(False)

补充:pytorch 中一个batch的训练过程

?
1
2
3
4
5
6
# 一般情况下
optimizer.zero_grad()             # 梯度清零
preds = model(inputs)             # inference,前向传播求出预测值
loss = criterion(preds, targets)  # 计算loss
loss.backward()                   # 反向传播求解梯度
optimizer.step()                  # 更新权重,更服务器之家络权重参数

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

以上为个人经验,希望能给大家一个参考,也希望大家多多支持服务器之家。

原文链接:https://blog.csdn.net/sinat_37532065/article/details/103990364

延伸 · 阅读

精彩推荐