Блог

Я изо всех сил пытаюсь создать генератор данных в PyTorch для извлечения 2D-изображений из множества 3D-кубов, сохраненных в .dat формате

Существует в общей сложности 200 3D-кубики, каждый из которых имеет определенную 128*128*128 форму. Теперь я хочу извлечь 2D-изображения из всех этих кубов по длине и ширине.

Например, a является ли куб размером 128*128*128

Поэтому я хочу извлечь все 2D-изображения по длине, т. Е., [:, i, :] что даст мне 128 2D-изображений по длине, и аналогично я хочу извлечь по ширине [:, :, i] , т. Е., что даст мне 128 2D-изображений по ширине. Поэтому я получаю в общей сложности 256 2D-изображений из 1 3D-куба, и я хочу повторить весь этот процесс для всех 200 кубов, предоставив мне 51200 2D-изображений.

До сих пор я пробовал очень простую реализацию, которая работает нормально, но для запуска требуется примерно 10 минут. Я хочу, чтобы вы, ребята, помогли мне создать более оптимальную реализацию с учетом сложности времени и пространства. Прямо сейчас мой текущий подход имеет временную сложность O(n2), можем ли мы продолжить его, чтобы уменьшить временную сложность

Я предоставляю ниже текущую реализацию

 from os.path import join as pjoin
import torch
import numpy as np
import os
from tqdm import tqdm
from torch.utils import data


class DataGenerator(data.Dataset):

    def __init__(self, is_transform=True, augmentations=None):

        self.is_transform = is_transform
        self.augmentations = augmentations
        self.dim = (128, 128, 128)

        seismicSections = [] #Input
        faultSections = [] #Ground Truth
        for fileName in tqdm(os.listdir(pjoin('train', 'seis')), total = len(os.listdir(pjoin('train', 'seis')))):
            unrolledVolSeismic = np.fromfile(pjoin('train', 'seis', fileName), dtype = np.single) #dat file contains unrolled cube, we need to reshape it
            reshapedVolSeismic = np.transpose(unrolledVolSeismic.reshape(self.dim)) #need to transpose the axis to get height axis at axis = 0, while length (axis = 1), and width(axis = 2)

            unrolledVolFault = np.fromfile(pjoin('train', 'fault', fileName),dtype=np.single)
            reshapedVolFault = np.transpose(unrolledVolFault.reshape(self.dim))

            for idx in range(reshapedVolSeismic.shape[2]):
                seismicSections.append(reshapedVolSeismic[:, :, idx])
                faultSections.append(reshapedVolFault[:, :, idx])

            for idx in range(reshapedVolSeismic.shape[1]):
                seismicSections.append(reshapedVolSeismic[:, idx, :])
                faultSections.append(reshapedVolFault[:, idx, :])

        self.seismicSections = seismicSections
        self.faultSections = faultSections

    def __len__(self):
        return len(self.seismicSections)

    def __getitem__(self, index):

        X = self.seismicSections[index]
        Y = self.faultSections[index]

        return X, Y
 

Пожалуйста, Помогите!!!

почему бы не сохранить только 3D-данные в mem и не позволить __getitem__ методу «нарезать» их на лету?

 class CachedVolumeDataset(Dataset):
  def __init__(self, ...):
    super(...)
    self._volumes_x = # a list of 200 128x128x128 volumes
    self._volumes_y = # a list of 200 128x128x128 volumes

  def __len__(self):
    return len(self._volumes_x) * (128   128)

  def __getitem__(self, index):
    # extract volume index from general index:
    vidx = index // (128   128)
    # extract slice index
    sidx = index % (128   128)
    if sidx < 128:
      # first dim
      x = self._volumes_x[vidx][:, :, sidx]
      y = self._volumes_y[vidx][:, :, sidx]
    else:
      sidx -= 128
      # second dim
      x = self._volumes_x[vidx][:, sidx, :]
      y = self._volumes_y[vidx][:, sidx, :]
    return torch.squeeze(x), torch.squeeze(y)