PyTorch 解决Dataset和Dataloader遇到的问题

yipeiwu_com5年前Python基础

今天在使用PyTorch中Dataset遇到了一个问题。先看代码

class psDataset(Dataset):
  def __init__(self, x, y, transforms = None):
    super(Dataset, self).__init__()
    self.x = x
    self.y = y
    if transforms == None:
      self.transforms = Compose([Resize((224, 224)), ToTensor()])
    else:
      self.transforms = transforms
    
  def __len__(self):
    return len(self.x)
  
  def __getitem__(self, idx):
    img = Image.open(self.x[idx])
    img = self.transforms(img)    
    return img, torch.tensor([[self.y[idx]]])

结果运行时报错:RuntimeError: invalid argument 0: Sizes of tensors must match except in dimension 0. Got 3 and 1 in dimension 1 at /opt/conda/conda-bld/pytorch_1522182087074/work/torch/lib/TH/generic/THTensorMath.c:2897

Google了一下发现是这样的:读入的图片有些是灰度图(1个通道),绝大多数是RGB图片(3通道),也有些是带透明度的(4通道)

。这导致在读入后最后一个维度(通道数)不一致(可能是1、3或者4)。

Dataloader在制作batch data时,tensor的shape必须一样,就报了这个错误。解决的方法是:img = img.convert(“RGB”)。完

整代码如下:

class psDataset(Dataset):
  def __init__(self, x, y, transforms = None):
    super(Dataset, self).__init__()
    self.x = x
    self.y = y
    if transforms == None:
      self.transforms = Compose([Resize((224, 224)), ToTensor()])
    else:
      self.transforms = transforms
    
  def __len__(self):
    return len(self.x)
  
  def __getitem__(self, idx):
    img = Image.open(self.x[idx])
    img = img.convert("RGB")
    img = self.transforms(img)    
    return img, torch.tensor([[self.y[idx]]])

以上这篇PyTorch 解决Dataset和Dataloader遇到的问题就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持【听图阁-专注于Python设计】。

相关文章

在Python的Django框架中编写错误提示页面

你应该在生产环境中把TEMPLATE_DEBUGFalse 如果这个设为`` True`` ,为了在那个好看的错误页面上显示足够的东西,Django的模版系统就会为每一个模版保存一些额外...

python 普通克里金(Kriging)法的实现

python 普通克里金(Kriging)法的实现

克里金法时一种用于空间插值的地学统计方法。 克里金法用半变异测定空间要素,要素即自相关要素。 半变异公式为: 其中γ(h) 是已知点 xi 和 xj 的半变异,***h***表示...

安装Python的web.py框架并从hello world开始编程

最近有一个小的web项目,想用喜爱都python,但是想到之前接触过都django我感觉一阵不寒而栗,为什么?Django的配置太过复杂,而且小项目不太适合MVC的开发模式,所以...

利用python实现冒泡排序算法实例代码

利用python实现冒泡排序算法实例代码

冒泡排序 冒泡排序(英语:Bubble Sort)是一种简单的排序算法。它重复地遍历要排序的数列,一次比较两个元素,如果他们的顺序错误就把他们交换过来。遍历数列的工作是重复地进行直到没...

Python实现简单登录验证

本文实例为大家分享了简单的Python登录验证,供大家参考,具体内容如下 编写登录接口 要求:1、输入用户名密码    2、认证成功后显示欢迎信息    3、输错三次后锁定 #cod...