今天在使用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遇到的問題就是小編分享給大家的全部內容了，希望能給大家一個參考，也希望大家多多支援我們。