Implemented train function in PyTorch - neural-network

Question

I'm building models of neural networks for some experiments. I use PyTorch and each time I train a model I use the following code:
def train_and_evaluate(net, optimizer, criterion):
start_time = time.time()
train_losses, test_losses, train_acc, test_acc = [], [], [], []
# net.double()
for epoch in range(num_epochs):
epoch_start_time = time.time()
running_loss_train, running_loss_test = 0.0, 0.0
total, correct, test_total, test_correct = 0, 0, 0, 0
# Train mode
net.train()
#Loop batches
for i, (images, labels) in enumerate(train_loader):
if use_cuda and torch.cuda.is_available():
images = images.cuda()
labels = labels.cuda()
optimizer.zero_grad()
outputs = net(images.detach())
# import pdb; pdb.set_trace()
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss_train += loss.item()
# Calculate training accuracy for epoch
_, predicted = torch.max(outputs.data, 1) # Get th
total += labels.size(0)
correct += (predicted == labels).sum().item()
# Calculate test loss and accuracy for epoch
net.eval()
with torch.no_grad():
for images, labels in test_loader:
if use_cuda and torch.cuda.is_available():
images = images.cuda()
labels = labels.cuda()
outputs = net(images)
test_loss = criterion(outputs, labels)
running_loss_test += test_loss.item()
_, predicted = torch.max(outputs.data, 1)
test_total += labels.size(0)
test_correct += (predicted == labels).sum().item()
if (epoch + 1) % 10 == 0 or epoch == 0:
print(f'Epoch [{epoch+1:02d}/{num_epochs}]\tTime: {time.time() - epoch_start_time:.2f}\tTrain Loss: {(running_loss_train / len(train_loader)):.4f}\tTrain Acc: {(correct / total):.0%}\tTest Loss: {running_loss_test / len(test_loader):.4f}\tTest Accuracy: {(test_correct / test_total):.0%}'.expandtabs(4))
train_losses.append(running_loss_train / len(train_loader))
test_losses.append(running_loss_test / len(test_loader))
train_acc.append(correct / total)
test_acc.append(test_correct / test_total)
print(f'Training time: {(time.time() - start_time)/60:5.2f} min.')
return train_losses, test_losses, train_acc, test_acc
I simply copy and paste it from other projects and I find it redundant to use that same code each time, it is useful almost in every model I test and usually looks quite similar. I wonder if there is a shortcut or some module with implementation of this function to reduce number of code lines and make it more readable. Something like:
model = NeuralNetwork()
train_losses, test_losses, train_acc, test_acc = train(model,
epochs=50,
verbose=True,
cuda=True,
optimizer=optimizer,
criterion=criterion)
That I could just pass the parameters for training and not explicitly 'write' the code evey time.

Answer 1

Use modules
Save train_and_evaluate function into python file. The file name is the module name with the suffix .py appended. for Ex.: model_util.py
Now you can access train_and_evaluate function using importing model_util module. like
import model_util
model = NeuralNetwork()
train_losses, test_losses, train_acc, test_acc = model_util.train_and_evaluate(...)