Как создать пользовательский загрузчик данных в pytorch для входной формы (242,242,12), где 12 - количество срезов трехмерного изображения МРТ Nifti

Question

Я пробую сеть двумерной свертки, и у меня возникли трудности при создании пользовательского загрузчика данных для моих входных изображений, которые представляют собой срезы (12) изображения МРТ размером (242 242).Я создал сеть, используя Keras, и она работает, однако она работает очень медленно, и я хотел бы изменить код на pytorch для производительности.

Вот мой загрузчик данных в Keras:

images=[]
labels =[]
res=(242,242)

for f in range(len(ff)):
a = nib.load(ff[f])
a = a.get_data()
a = a[:,80:103:2,:]
a =a.transpose(0,2,1)
a = resize(a,res)

for i in range(a.shape[2]):
  if('A' in os.path.basename(ff[f])):
    images.append((a[:,:,i]))
    label=0
    labels.append(label)
  else:
    images.append((a[:,:,i]))
    label=1
    labels.append(label)

#Convert labels and images as arrays and getting categorical labels

from sklearn.model_selection import train_test_split
train_X,valid_X,train_ground,valid_ground = train_test_split(images,
                                                         labels,
                                                         test_size=0.2,
                                                         random_state=1)
def model(images):
# my conv net defined

model_train = model.fit(train_X, train_ground, batch_size=batch_size,epochs=epochs,verbose=1,validation_data=(valid_X, valid_ground))

Как мне нужно изменить этот код в Pytorch, чтобы получить загрузчик данных для обучения.

Любая помощь будет принята с благодарностью.

Спасибо