I’ve been working on training RNNs with PyTorch for some time now, and I’m wondering about the optimal architecture for doing so, using all the features that are available (padding, packing, etc.). I haven’t really found any tutorial or documentation that would provide a general framework of doing things properly from the beginning to the end.
The objectives are the following:
- PyTorch-heavy (using all the natively supported features)
- parallel DataLoader
- CUDA-heavy (put as much of the work as possible to GPU)
- computationally optimal (as a result of all the above)
The assumptions of the problem I’m solving (although I believe not crucial) are the following:
- regression problem
- training on all signals
- predicting/evaluating last signal
I’ve come up with the following architecture
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import *
from torch.utils.data import Dataset, DataLoader
class RNNDataset(Dataset):
def __init__(self, num_sequences, num_features, test=False):
# random lengths from 1 to 100
random_sequence_lengths = torch.randint(1, 100, (num_sequences,))
# List of tensors, 1 element representing one sequence
self.train = [torch.rand(l, num_features) for l in random_sequence_lengths]
# List of sequence lengths
self.lens = random_sequence_lengths
self.len = len(self.train)
self.test = test
if test:
# List of tensors, 1 element representing labels for one sequence
self.labels = [torch.rand(l, 1) for l in random_sequence_lengths]
else:
self.labels = None
def __getitem__(self, idx):
if self.test:
labels = self.labels[idx]
else:
labels = None
return self.train[idx], self.lens[idx], labels
def __len__(self):
return self.len
class RNNNet(torch.nn.Module):
def __init__(self, num_features):
super(RNNNet, self).__init__()
self.pre_rnn = nn.Linear(num_features, num_features * 2)
self.rnn = nn.GRU(num_features * 2, num_features, batch_first=True)
self.post_rnn = nn.Linear(num_features, 1)
def forward(self, inputs, lens):
# Indices for proper input order with regard to labels
indices = torch.argsort(-lens)
indices_back = torch.argsort(indices)
# Pre RNN FC layer
inputs_fc_pre = F.relu(self.pre_rnn(inputs))
# Pack the sequence
inputs_packed = pack_padded_sequence(inputs_fc_pre[indices], lens[indices], batch_first=True)
# RNN layer
inputs_rnn, _ = self.rnn(inputs_packed)
# Reverse operation, pad the packed sequence
inputs_rnn_padded, _ = pad_packed_sequence(inputs_rnn, batch_first=True)
# Post RNN FC layer
input_post_rnn = self.post_rnn(F.relu(inputs_rnn_padded))[indices_back]
return input_post_rnn
def predict(self, test, validation=False):
# Predicting only the last output
# Depending on whether it's a monitoring phase or an actual prediction, different preparation is needed
if type(test) is DataLoader:
test_loader = test
else:
test_loader = DataLoader(test, batch_size=100, shuffle=False,
num_workers=2, collate_fn=rnn_collate, pin_memory=True)
self.eval()
device = next(model.parameters()).device
results = []
with torch.no_grad():
for inputs_part, lens_part, labels in test_loader:
results_chunk = self.forward(
inputs_part.to(device, non_blocking=True),
lens_part.to(device, non_blocking=True)
)
# Getting last output from test prediction
last_preds = [x[l - 1] for x, l in zip(results_chunk, lens_part)]
if validation:
last_labels = [x[l - 1] for x, l in zip(labels, lens_part)]
results.extend(zip(last_preds, last_labels))
else:
results.extend(last_preds)
results_final = torch.Tensor(results).detach()
self.train()
return results_final
def rnn_collate(batch):
inputs, lens, labels = zip(*batch)
# Pad the sequences and labels to equal length.
# If processing test set, all labels will be None, in that case no processing is done.
if labels[0] is not None:
labels = pad_sequence(labels, batch_first=True, padding_value=-999)
else:
labels = None
inputs = pad_sequence(inputs, batch_first=True, padding_value=0)
lens = torch.LongTensor(lens)
return inputs, lens, labels
def fit(train, validation, model, num_epochs, batch_size, learning_rate, device):
criterion = nn.MSELoss()
train_loader = DataLoader(train, batch_size=batch_size, shuffle=True,
num_workers=2, collate_fn=rnn_collate, pin_memory=True)
valid_loader = DataLoader(valid, batch_size=batch_size, shuffle=False,
num_workers=2, collate_fn=rnn_collate, pin_memory=True)
model.zero_grad()
model.train()
model.to(device, non_blocking=True)
optimizer = torch.optim.Adam(model.parameters(), learning_rate)
for epoch in range(num_epochs):
running_loss = 0.0
for i, data in enumerate(train_loader):
# Get the inputs
inputs, lens, labels = data
labels = labels.to(device, non_blocking=True)
inputs = inputs.to(device, non_blocking=True)
lens = lens.to(device, non_blocking=True)
# Zero the parameter gradients
optimizer.zero_grad()
# Forward + backward + optimize
outputs = model.forward(inputs, lens)
# Learning on all time steps outputs - need a mask to extract the relevant ones
labels_mask = (labels != -999)
# Extract outputs of interest
outputs_extracted = outputs[labels_mask]
# Get the RMSE
loss = torch.sqrt(criterion(outputs_extracted, labels[labels_mask]))
loss.backward()
optimizer.step()
# Gather statistics
running_loss += loss.item()
results = model.predict(valid, True)
test_loss = torch.sqrt(F.mse_loss(results[:, 0], results[:, 1])).item()
print(f"Train loss after epoch {epoch}: {running_loss / i}")
print(f"Test loss after epoch {epoch}: {test_loss}")
# Reset the loss monitor
running_loss = 0.0
return model
if __name__ == "__main__":
num_sequences = 1000
num_features = 50
device = "cuda"
num_epochs = 10
batch_size = 64
learning_rate = 1e-3
train = RNNDataset(num_sequences, num_features, test=True)
valid = RNNDataset(num_sequences, num_features, test=True)
test = RNNDataset(num_sequences, num_features, test=False)
model = RNNNet(num_features)
model = fit(train, valid, model, num_epochs, batch_size, learning_rate, device)
preds = model.predict(test)
print(preds[:10])
It works pretty nicely and fast in general. The main problem being that padding of the batch is done on CPU and not on GPU in rnn_collate. I tried sending the batch elements to GPU in rnn_collate or directly in RNNDataset (changing the multiprocessing to spawn to avoid CUDA initialization errors) but I was getting some shared memory block exceptions (with memory shared block already increased), so I assumed the approach is not optimal.
My questions are the following:
- Is there some tutorial/documentation on how to create an optimal (or close to) RNN training architecture in PyTorch?
- Do you have any hints regarding the above code that would make it faster/more optimal/more PyTorch-heavy?
All suggestions welcome!