如何在 PyTorch 中有效地制作一小批量图像



我正在尝试使用pytorch中预先训练的ResNet模型计算前向传递。我在创建小批量的 4-D 张量时遇到问题。有人可以告诉正确的方法是什么吗?

编辑:我更改了代码,它现在可以工作了。但是,我仍然认为应该有一种更有效的方法。

这是我的代码:

import pickle
import json
import shutil
import Image
import torchvision.models as models
import torchvision.transformers as transformers
from torch.autograd import Variable
from torch import Tensor
import glob
import torch
batch_size = 128
im_size = 299
normalize = transforms.Normalize(
   mean=[0.485, 0.456, 0.406],
   std=[0.229, 0.224, 0.225]
)
preprocess = transforms.Compose([
   transforms.Scale(im_size),
   transforms.CenterCrop(im_size),
   transforms.ToTensor(),
   normalize
])

model = models.resnet50(pretrained=True)
d_batch = make_batch(imgs, batch_size)
dtype = torch.FloatTensor
tmp = Variable(torch.randn(batch_size, 3, im_size, im_size).type(dtype), requires_grad=False)

for batch in tqdm(batches):
        try:
                data = [Image.open(img) for img in batch]
                for idx, item in enumerate(data):
                        tmp[idx] = preprocess(item)
                batch_result = model(tmp)
        except Exception,x:
                print x

使用dataset = torchvision.datasets.ImageFolder(...)可以从图像文件夹加载数据集。之后,您可以使用torch.utils.data.DataLoader(dataset, batch_size=batchSize)指定小批量大小和其他内容以进行进一步处理。

相关内容

  • 没有找到相关文章

最新更新