Metrics not being logged properly on remote GPU
aaronwtr opened this issue · comments
🐛 Bug
I have written some PyTorch Lightning code for which I am using torchmetrics to evaluate. Locally on my CPU everything works fine. However, when I move my script to a cluster, and try to run it on a single GPU, I run into problems. Specifically, it seems like my error metrics are not being calculated properly (see attached image). What could be the cause here?
To Reproduce
I am loading my data with the LinkNeighbourLoader provided by PyG. I suspect the observed behaviour might be partially attributable to the way LinkNeighbourLoader sends (or doesn't send?) the data to the correct device. Torchmetrics should handle this out-of-the-box, but it might be that LinkNeighbourLoader doesn't integrate properly here.
I load my data as follows:
class HGSLNetDataModule(pl.LightningDataModule):
def __init__(self, train_graph, test_graph, config):
super().__init__()
self.config = config
self.train_graph = train_graph
self.val_graph = None
self.test_graph = test_graph
self.scaler = None
def setup(self, stage=None):
if stage == 'fit' or stage is None:
self.train_graph, self.val_graph, _ = self.link_split_transform()
if os.path.exists('cache/quantile_scaler.pkl'):
with open('cache/quantile_scaler.pkl', 'rb') as f:
self.scaler = pkl.load(f)
self.train_graph.x = self.scaler.transform(self.train_graph.x)
else:
self.scaler = QuantileTransformer()
self.train_graph.x = self.scaler.fit_transform(self.train_graph.x)
with open('cache/quantile_scaler.pkl', 'wb') as f:
pkl.dump(self.scaler, f)
self.val_graph.x = self.scaler.transform(self.val_graph.x)
self.train_graph.x = torch.from_numpy(self.train_graph.x).to(torch.float32)
self.val_graph.x = torch.from_numpy(self.val_graph.x).to(torch.float32)
self.train_edge_labels, self.train_edge_label_index = (self.train_graph.edge_label,
self.train_graph.edge_label_index)
self.val_edge_labels, self.val_edge_label_index = self.val_graph.edge_label, self.val_graph.edge_label_index
elif stage == 'test':
self.test_graph = self.test_graph
with open('cache/quantile_scaler.pkl', 'rb') as f:
self.scaler = pkl.load(f)
self.test_graph.x = self.scaler.transform(self.test_graph.x)
self.test_graph.x = torch.from_numpy(self.test_graph.x).to(torch.float32)
self.test_edge_labels, self.test_edge_label_index = (self.test_graph.edge_label,
self.test_graph.edge_index)
def train_dataloader(self):
return LinkNeighborLoader(
self.train_graph,
batch_size=self.config['batch_size'],
edge_label=self.train_edge_labels,
edge_label_index=self.train_edge_label_index,
num_neighbors=[20, 10],
pin_memory=True,
shuffle=True
)With similar loaders for my test and val data. Then, I initialise my metrics in where I am defining the Lightning training and testing logic as follows:
class HGSLNet(pl.LightningModule):
def __init__(self, num_layers, hidden_channels, num_heads, config):
super().__init__()
self.config = config
self.model = GATModel(num_layers, hidden_channels, num_heads)
self.best_avg_mcc = 0
self.best_avg_acc = 0
self.val_mcc = MatthewsCorrCoef(task='binary')
self.val_acc = Accuracy(task='binary')
self.test_mcc = MatthewsCorrCoef(task='binary')
self.test_acc = Accuracy(task='binary')E.g., in my validation step, I log the metrics as follows:
def validation_step(self, batch, batch_idx):
x, edge_label_index, y = batch.x, batch.edge_label_index, batch.edge_label
logit, proba, pred = self(x, edge_label_index)
_y = y.float()
self.val_batch_size = batch.edge_label_index.size(1)
loss = F.binary_cross_entropy_with_logits(logit, _y)
val_mcc = self.val_mcc(pred, y)
val_acc = self.val_acc(pred, y)
self.log('val_loss', loss, on_step=False, on_epoch=True, batch_size=self.val_batch_size)
self.log('val_mcc', val_mcc, on_step=False, on_epoch=True, batch_size=self.val_batch_size)
self.log('val_acc', val_acc, on_step=False, on_epoch=True, batch_size=self.val_batch_size)
return lossThe logic is similar for my training and testing steps. This leads to the loss curves as displayed in the attached figure, where green is the run on HPC cluster GPUs and red locally on CPU.
Expected behavior
The metrics to be logged at each epoch properly.
Environment
Python environment:
- Python==3.10.7
- PyTorch==2.2.0+cu118
- cudnn==8.9.4
- cuda==12.2
- gcc==12.1.0
- lightning==2.1.3
- torchmetrics== 1.3.0.post0
- torch_geometric==2.4.0
- pyg-lib==0.4.0
OS:
-CentOS Linux 7 (Core)
Additional context
Loss curves:
Hi! thanks for your contribution!, great first issue!
Hi @aaronwtr, thanks for raising this issue.
Have you solved the issue or do it still persist? I think a bit more information is needed here. The logging behavior should really not change when you move from one device to another (even in this case another computer system). I wonder if it is the logging that is going wrong here or is it the actually training that is going wrong e.g. if you just print to the terminal do you still produce constant metric values?
Additionally, if you somehow can provide a fully reproducible example that would be nice