[Python] Cross-validation with PyTorch

I am trying to implement cross-validation with PyTorch;

I am used to Sklearn's cross_validation function which is much simpler to use (most likely because I did not fully understand what was happening under the hood); I was hoping someone could advise if first of all this looks correct and, if it does, if there is a better way to write this.

I have looked at other posts, and got inspiration from there but I could not find exactly what I am trying to create, which is:

1. Loop through the dataset for n_splits (in my case 5).
2. On each loop, train the model on the train split.
3. Evaluate the model on the test split.
4. Calculate the RMSE of the model evaluation on the test split.
5. At the end of the 5 folds, return the average RMSE of the 5 folds.

And as I was creating this, I have started wondering: how should one decide the number of epochs for cross-validation? I always assumed that this would just be 1 but from what I saw in other posts this number can actually vary.

Here is my code:

X = main_df.drop(columns=['target'])
y = main_df['target']

X_train_t = torch.tensor(X.values, dtype=torch.float32)
y_train_t = torch.tensor(y.values, dtype=torch.float32)

cv_dataset = TensorDataset(X_train_t, y_train_t)


class NetModel(nn.Module):
def __init__(self, in_count):
super(NetModel, self).__init__()
self.layers = nn.Sequential(
nn.Linear(in_count, 32),
nn.ReLU(),
nn.Linear(32, 1)
)

def forward(self, x):
out = self.layers(x)
return out


num_epochs = 1
kfold = KFold(n_splits=5, shuffle=True, random_state=42)
tot_rmse = 0.0

for fold, (train_ids, test_ids) in enumerate(kfold.split(cv_dataset)):

print(f'FOLD {fold+1}')

train_subsampler = torch.utils.data.SubsetRandomSampler(train_ids)
test_subsampler = torch.utils.data.SubsetRandomSampler(test_ids)

train_loader = DataLoader(cv_dataset, batch_size=5, sampler=train_subsampler)
test_loader = DataLoader(cv_dataset, batch_size=5, sampler=test_subsampler)

# initialize model
model = NetModel(X_train_t.shape[1])
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.0001)

for epoch in range(num_epochs):
current_loss = 0.0

model.train()
for i, data in enumerate(train_loader, 0):

inputs, targets = data
targets = targets.reshape((targets.shape[0], 1))

optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()

fold_rmse = 0.0
model.eval()
with torch.no_grad():
for i, data in enumerate(test_loader, 0):

inputs, targets = data
targets = targets.reshape((targets.shape[0], 1))

outputs = model(inputs)

rmse = root_mean_squared_error(targets, outputs)

# accumulate rmse of the fold
fold_rmse += rmse

print(f"Sum of fold {fold+1} RMSE: {fold_rmse:.3f} | Count batches: {i+1}")

# calculate average of the fold rmse (accumulated rmse / number of batches)
tot_rmse += fold_rmse/(i+1)
print("--------------------------------------------------")
print(f"Avg. RMSE: {tot_rmse/5}")


This is the output of the above:

FOLD 1
Sum of fold 1 RMSE: 2124.690 | Count batches: 89
--------------------------------------------------
FOLD 2
Sum of fold 2 RMSE: 1817.517 | Count batches: 89
--------------------------------------------------
FOLD 3
Sum of fold 3 RMSE: 2160.074 | Count batches: 89
--------------------------------------------------
FOLD 4
Sum of fold 4 RMSE: 1613.786 | Count batches: 89
--------------------------------------------------
FOLD 5
Sum of fold 5 RMSE: 1713.522 | Count batches: 89
--------------------------------------------------
Avg. RMSE: 21.190088516406796


Thank you!

Continue reading...