如何在PyTorch Lightning中从prepare_data()获取数据集到setup()

Question

我使用 PyTorch Lightning 的

NumPy

方法在

prepare_data()

方法中使用

DataModules

制作了自己的数据集。现在，我想将数据传递到

setup()

方法中以分为训练和验证。

import numpy as np 
import pytorch_lightning as pl 
from torch.utils.data import random_split, DataLoader, TensorDataset
import torch
from torch.autograd import Variable
from torchvision import transforms

np.random.seed(42)

device = 'cuda' if torch.cuda.is_available() else 'cpu'

class DataModuleClass(pl.LightningDataModule):
    def __init__(self):
        super().__init__()
        self.constant = 2
        self.batch_size = 10
        
    def prepare_data(self):
        a = np.random.uniform(0, 500, 500)
        b = np.random.normal(0, self.constant, len(a))

        c = a + b
        X = np.transpose(np.array([a, b]))
        
        # Converting numpy array to Tensor
        self.x_train_tensor = torch.from_numpy(X).float().to(device)
        self.y_train_tensor = torch.from_numpy(c).float().to(device)
        
        training_dataset = TensorDataset(self.x_train_tensor, self.y_train_tensor)

        return training_dataset
    
    def setup(self):
        data = # What I have to write to get the data from prepare_data()
        self.train_data, self.val_data = random_split(data, [400, 100])
        
        
    def train_dataloader(self):
        training_dataloader = setup() # Need to get the training data
        return DataLoader(self.training_dataloader)

    def val_dataloader(self):
        validation_dataloader = prepare_data() # Need to get the validation data
        return DataLoader(self.validation_dataloader)
    
obj = DataModuleClass()
print(obj.prepare_data())  

Answer 1

与您之前的问题相同的答案...

def prepare_data(self):
    a = np.random.uniform(0, 500, 500)
    b = np.random.normal(0, self.constant, len(a))

    c = a + b
    X = np.transpose(np.array([a, b]))

    # Converting numpy array to Tensor
    self.x_train_tensor = torch.from_numpy(X).float().to(device)
    self.y_train_tensor = torch.from_numpy(c).float().to(device)

    training_dataset = TensorDataset(self.x_train_tensor, self.y_train_tensor)

    self.training_dataset = training_dataset

def setup(self):
    data = self.training_dataset
    self.train_data, self.val_data = random_split(data, [400, 100])
    
    
def train_dataloader(self):
    return DataLoader(self.train_data)

def val_dataloader(self):
    return DataLoader(self.val_data)

Answer 2

在 DataModule 对象上调用

setup()

后，只需在该对象上调用

prepare()

即可。所以：

dm = DataModuleClass()
dm.prepare_data()
dm.setup()

Answer 3

prepare_data()

仅在主进程中用于数据下载或其他数据操作。如果你想制作自己的数据集或从磁盘读取数据，应该在

setup()

中完成，它支持多处理。

对于您的问题，您应该删除

prepare_data()

并将其代码合并到

setup()

中以拆分数据集。