擅长:python、mysql、java
<p>因为您有图像,所以可能需要对它们执行转换。所以<code>TensorDataset</code>不是这里最好的选择。相反,您可以创建自己的<code>Dataset</code>。像这样的:</p>
<pre><code>import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np
from PIL import Image
class MyDataset(Dataset):
def __init__(self, data, targets, transform=None):
self.data = data
self.targets = torch.LongTensor(targets)
self.transform = transform
def __getitem__(self, index):
x = self.data[index]
y = self.targets[index]
if self.transform:
x = Image.fromarray(self.data[index].astype(np.uint8).transpose(1,2,0))
x = self.transform(x)
return x, y
def __len__(self):
return len(self.data)
# Let's create 10 RGB images of size 128x128 and ten labels {0, 1}
data = list(np.random.randint(0, 255, size=(10, 3, 128, 128)))
targets = list(np.random.randint(2, size=(10)))
transform = transforms.Compose([transforms.Resize(64), transforms.ToTensor()])
dataset = MyDataset(data, targets, transform=transform)
dataloader = DataLoader(dataset, batch_size=5)
</code></pre>