I’m new to PyTorch, and I was wondering if there’s any implementation of (or a reason why I wouldn’t want) an extension of nn.Module which contains a method that trains the model? I was thinking something like:
class TrainModule(nn.Module):
def __init__(self):
super().__init__()
def train_model(self, traindata, valdata, epochs, batch_size,
criterion, optimizer, scheduler):
from tqdm import tqdm
trainloader = DataLoader(traindata, shuffle = True,
batch_size = batch_size)
valloader = DataLoader(valdata, shuffle = True,
batch_size = batch_size)
num_batches = len(traindata) // batch_size
if len(traindata) % batch_size:
num_batches += 1
for epoch in range(epochs):
with tqdm(total = len(traindata)) as epoch_pbar:
epoch_pbar.set_description(f'Epoch {epoch}')
acc_loss = 0
for idx, (inputs, target) in enumerate(trainloader):
# Reset the gradients
optimizer.zero_grad()
# Do a forward pass, calculate loss and backpropagate
outputs = self.forward(inputs)
loss = criterion(outputs, target)
loss.backward()
optimizer.step()
# Add loss to accumulated loss
acc_loss += loss
# Update progress bar description
avg_loss = acc_loss / (idx + 1)
desc = f'Epoch {epoch} - loss {avg_loss:.4f}'
epoch_pbar.set_description(desc)
epoch_pbar.update(inputs.shape[0])
# Learning rate decay
avg_loss = acc_loss / num_batches
scheduler.step(avg_loss)
# Deal with validation data
(...)
And allowing string inputs such as ‘adam’ for optimizer, or ‘mse’ for criterion, and so on.
Not having to write out these every time I build a model would be great, and it’d be consistent with the current use, as you can simply not use the train_model method, or overwrite it.
