Как получить имя файла образца из DataLoader? - PullRequest
Новая
       26

Как получить имя файла образца из DataLoader?

3 голосов
/ 21 июня 2019

Мне нужно написать файл с результатом теста данных сверточной нейронной сети, которую я обучил. Данные включают сбор речевых данных. Формат файла должен быть «имя файла, прогноз», но мне трудно извлечь имя файла. Я загружаю данные так: 

import torchvision
from torchvision import transforms
from torch.utils.data import DataLoader

TEST_DATA_PATH = ...

trans = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

test_dataset = torchvision.datasets.MNIST(
    root=TEST_DATA_PATH,
    train=False,
    transform=trans,
    download=True
)

test_loader = DataLoader(dataset=test_dataset, batch_size=1, shuffle=False)

и я пытаюсь записать в файл следующее: 

f = open("test_y", "w")
with torch.no_grad():
    for i, (images, labels) in enumerate(test_loader, 0):
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        file = os.listdir(TEST_DATA_PATH + "/all")[i]
        format = file + ", " + str(predicted.item()) + '\n'
        f.write(format)
f.close()

Проблема с os.listdir(TESTH_DATA_PATH + "/all")[i] заключается в том, что он не синхронизирован с порядком загруженных файлов test_loader. Что я могу сделать?

Ответы [ 2 ]

0 голосов
/ 21 июня 2019

В общем случае DataLoader предоставит вам партии из набора (ов) данных, которые у него есть.

AS @Barriel, упомянутый в случае проблем классификации с одной / несколькими метками, DataLoader не имеет имени файла изображения, только тензоры, представляющие изображения, и классы / метки.

Тем не менее, конструктор DataLoader при загрузке объектов может иметь небольшие значения (вместе с набором данных вы можете упаковать цели / метки и имена файлов, если хотите)

Таким образом, DataLoader может каким-то образом захватить то, что вам нужно.

0 голосов
/ 21 июня 2019

Ну, это зависит от того, как реализован ваш Dataset.Например, в случае torchvision.datasets.MNIST(...) вы не можете получить имя файла просто потому, что не существует такой вещи, как имя файла отдельной выборки (выборки MNIST загружаются другим способом ).

Поскольку вы не показали свою реализацию Dataset, я расскажу вам, как это можно сделать с помощью torchvision.datasets.ImageFolder(...) (или любого torchvision.datasets.DatasetFolder(...)): 

f = open("test_y", "w")
with torch.no_grad():
    for i, (images, labels) in enumerate(test_loader, 0):
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        sample_fname, _ = test_loader.dataset.samples[i]
        f.write("{}, {}\n".format(sample_fname, predicted.item()))
f.close()

Вы можете видеть, что путь к файлу извлекается во время __getitem__(self, index), в частности здесь .

Если вы реализовали свой собственный Dataset (и, возможно,хотел бы поддерживать shuffle и batch_size > 1), тогда я бы вернул sample_fname при вызове __getitem__(...) и сделал бы что-то вроде этого: 

for i, (images, labels, sample_fname) in enumerate(test_loader, 0):
    # [...]

Таким образом, вам не нужно было бызаботиться о shuffle.И если batch_size больше 1, вам нужно изменить содержимое цикла на что-то более общее, например: 

f = open("test_y", "w")
for i, (images, labels, samples_fname) in enumerate(test_loader, 0):
    outputs = model(images)
    pred = torch.max(outputs, 1)[1]
    f.write("\n".join([
        ", ".join(x)
        for x in zip(map(str, pred.cpu().tolist()), samples_fname)
    ]) + "\n")
f.close()
...