加入收藏 | 设为首页 | 会员中心 | 我要投稿 李大同 (https://www.lidatong.com.cn/)- 科技、建站、经验、云计算、5G、大数据,站长网!
当前位置: 首页 > 编程开发 > Python > 正文

PyTorch 解决Dataset和Dataloader遇到的问题

发布时间:2020-12-17 17:33:46 所属栏目: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))

今天在使用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,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遇到的问题就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持我们。

(编辑:李大同)

【声明】本站内容均来自网络,其相关言论仅代表作者个人观点,不代表本站立场。若无意侵犯到您的权利,请及时与联系站长删除相关内容!

    推荐文章
      热点阅读