正如@jodag所建议的那样,使用DataLoaders是一个好主意。
我有一段代码片段在Pytorch中用于我的一些CNN。
from torch.utils.data import Dataset, DataLoader
import torch
class Data(Dataset):
"""
Constructs a Dataset to be parsed into a DataLoader
"""
def __init__(self,X,y):
X = torch.from_numpy(X).float()
X = torch.transpose(X,1,2)
y = torch.from_numpy(y).float()
self.X,self.y = X,y
def __getitem__(self, i):
return self.X[i],self.y[i]
def __len__(self):
return self.X.shape[0]
def create_data_loader(X,y,batch_size,**kwargs):
"""
Creates a data-loader for the data X and y
params:
-------
X: np.array
- numpy array of size "n" x k where n is samples an "k" is number of features
y: np.array
- numpy array of sie "n"
batch_size: int
- Take a wild guess, dumbass
kwargs:
- Additional keyword-arguments for "DataLoader"
return
------
dl: torch.utils.data.DataLoader object
"""
data = Data(X, y)
dl = DataLoader(data, batch_size=batch_size,num_workers=0,**kwargs)
return dl
这是这样使用的:
from create_data_loader import create_data_loader
train_data_loader= create_data_loader(X_train,y_train,batch_size=32)
val_data_loader= create_data_loader(X_val,y_val,batch_size=32,shuffle=False)
for x_train, y_train in train_data_loader:
logit = net(x_train,y_train)
.
.
net.eval()
for x_val,y_val in val_data_loader:
logit = net(x_val,y_val)
classes_pred = logit.argmax(axis=1)
print(f"Val accuracy: {(y_val==classes_pred).mean()}")
torch.utils.data.DataLoader
似乎更合适,而不是编写自己的生成器。 - jodag