如何更改PyTorch中的DataLoader以读取一个图像进行预测

Question

目前，我有一个预训练的模型，它使用DataLoader读取一批图像来训练模型。

self.data_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False, 
num_workers=1, pin_memory=True)
...
model.eval()
for step, inputs in enumerate(test_loader.data_loader):
outputs = model(torch.cat([inputs], 1))
...

当图像从队列中到达时，我想对它们进行处理(做出预测(。它应该类似于读取单个图像并运行模型对其进行预测的代码

from PIL import Image
new_input = Image.open(image_path)
model.eval()
outputs = model(torch.cat([new_input ], 1))

我想知道您是否可以指导我如何做到这一点，并在DataLoader中应用相同的转换。

Answer 1

您可以使用IterableDataset:

from torch.utils.data import IterableDataset
class MyDataset(IterableDataset):
def __init__(self, image_queue):
self.queue = image_queue
def read_next_image(self):
while self.queue.qsize() > 0:
# you can add transform here
yield self.queue.get()
return None
def __iter__(self):
return self.read_next_image()

和batch_size=1:

import queue
import torchvision.transforms.functional as TF
buffer = queue.Queue()
new_input = Image.open(image_path)
buffer.put(TF.to_tensor(new_input)) 
# ... Populate queue here
dataset = MyDataset(buffer)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=1)
for data in dataloader:
model(data) # data is one-image batch of size [1,3,H,W] where 3 - number of color channels

Answer 2

我不知道dataLoader，但您可以使用以下函数加载单个图像：

def safe_pil_loader(path, from_memory=False):
try:
if from_memory:
img = Image.open(path)
res = img.convert('RGB')
else:
with open(path, 'rb') as f:
img = Image.open(f)
res = img.convert('RGB')
except:
res = Image.new('RGB', (227, 227), color=0)
return res

对于应用转换，您可以执行以下操作：

trans = transforms.Compose([
transforms.Resize(299),
transforms.CenterCrop(299),
transforms.ToTensor(),
normalize,
])
img=trans(img)