PyTorch 解决Dataset和Dataloader遇到的问题

yipeiwu_com6年前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 基础之字符串string详解及实例

Python字符串(string) 详解 及 代码 Python的字符串可以使用单引号('), 双引号("), 三引号('''); 三引号(''')里面, 可以添加单引号和双引号, 也可...

python 调用HBase的简单实例

新来的一个工程师不懂HBase,java不熟,python还行,我建议他那可以考虑用HBase的thrift调用,完成目前的工作。 首先,安装thrift 下载thrift,这里,我用的...

python用户管理系统

本文实例为大家分享了Python用户管理系统的具体代码,供大家参考,具体内容如下 用户管理系统 1.注册新用户 如果注册用户已经存在,则报错 需要填写信息: name, passw...

python多进程提取处理大量文本的关键词方法

python多进程提取处理大量文本的关键词方法

经常需要通过python代码来提取文本的关键词,用于文本分析。而实际应用中文本量又是大量的数据,如果使用单进程的话,效率会比较低,因此可以考虑使用多进程。 python的多进程只需要使用...

python循环定时中断执行某一段程序的实例

问题说明 最近在写爬虫,由于单个账号访问频率太高会被封,所以需要在爬虫执行一段时间间隔后自己循环切换账号 所以就在想,有没有像单片机那样子设置一个定时中断,再定义一个中断入口,这样子每隔...