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内置模块hashlib、hmac与uuid用法分析

本文实例讲述了Python内置模块hashlib、hmac与uuid用法。分享给大家供大家参考,具体如下: 一、hashlib md5和sha算法通过消息摘要算法生成定长的消息摘要,消息...

Python 网络编程起步(Socket发送消息)

Python 网络编程起步(Socket发送消息)

一、服务端(Server.py)    服务端要做的事情是:    1. 创建一个Socket对象 ˂!-- Code hi...

python安装pil库方法及代码

python安装pil库方法及代码

安装PIL 在Debian/Ubuntu Linux下直接通过apt安装: $ sudo apt-get install python-imaging Mac和其他版本的Linux...

scrapy-redis源码分析之发送POST请求详解

1 引言 这段时间在研究美团爬虫,用的是scrapy-redis分布式爬虫框架,奈何scrapy-redis与scrapy框架不同,默认只发送GET请求,换句话说,不能直接发送POST...

在cmd命令行里进入和退出Python程序的方法

在cmd命令行里进入和退出Python程序的方法

进入: 直接输入python即可,如图所示 退出: 1:输入exit(),回车 2:输入quit(),回车 3:输入ctrl+z,回车 以上这篇在cmd命令行里进入和退出Python程...