python - PyTorch vs MXNet: Significant Performance Discrepancy in Graph Neural Network Implementation - Stack Overflow

admin管理员组

文章数量:1122832

I am converting a Graph Neural Network (GNN) model from MXNet to PyTorch and experiencing significant performance discrepancies during training. The same architecture and data yield much worse results in PyTorch compared to the MXNet implementation. Github implementaion at:

The model is a Spatio-Temporal Graph Convolutional Network (ST-GNN) used for traffic forecasting. The input data has the shape (num_samples, time_steps, num_nodes, num_features). The adjacency matrix (adj_mx) is constructed using a combination of connectivity and dynamic time warping (DTW) matrices. Training Results: MXNet Results (First 3 Epochs):

training: Epoch: 1, RMSE: 129.24, MAE: 96.25, time: 128.01 s
validation: Epoch: 1, loss: 40.85, time: 144.13 s
test: Epoch: 1, MAE: 39.81, MAPE: 29.27, RMSE: 58.48, time: 165.47s

training: Epoch: 2, RMSE: 50.62, MAE: 34.54, time: 124.97 s
validation: Epoch: 2, loss: 32.37, time: 140.61 s
test: Epoch: 2, MAE: 31.57, MAPE: 23.59, RMSE: 45.88, time: 161.53s

training: Epoch: 3, RMSE: 43.97, MAE: 29.62, time: 125.43 s
validation: Epoch: 3, loss: 31.18, time: 141.05 s
test: Epoch: 3, MAE: 30.52, MAPE: 24.56, RMSE: 44.58, time: 161.73s

PyTorch Results (First 3 Epochs):

Training: Epoch: 1, Loss: 225979.7226, Time: 41.94s
Validation: Epoch: 1, Loss: 222.6407, Time: 46.58s
Test: Epoch: 1, MAE: 219.72, MAPE: 101.28, RMSE: 270.32, Time: 53.61s

Training: Epoch: 2, Loss: 62675.5868, Time: 42.25s
Validation: Epoch: 2, Loss: 155.2467, Time: 47.07s
Test: Epoch: 2, MAE: 154.06, MAPE: 91.94, RMSE: 206.60, Time: 55.12s

Training: Epoch: 3, Loss: 30421.8932, Time: 42.49s
Validation: Epoch: 3, Loss: 130.6325, Time: 47.29s
Test: Epoch: 3, MAE: 127.57, MAPE: 147.89, RMSE: 166.99, Time: 55.14s

Config:

{
    "module_type": "individual",
    "act_type": "GLU",
    "temporal_emb": true,
    "spatial_emb": true,
    "use_mask": true,
    "first_layer_embedding_size": 64,
    "filters": [
        [64, 64, 64],
        [64, 64, 64],
        [64, 64, 64]
    ],
    "batch_size": 32,
    "optimizer": "adam",
    "learning_rate": 1e-3,
    "epochs": 200,
    "max_update_factor": 1,
    "ctx": 0,
    "adj_filename": "./data/PEMS04/PEMS04.csv",
    "id_filename": null,
    "graph_signal_matrix_filename": "./data/PEMS04/PEMS04.npz",
    "num_of_vertices": 307,
    "points_per_hour": 12,
    "num_for_predict": 12,
    "num_of_features": 1,
    "adj_dtw_filename": "./data/adj_PEMS04_001.csv"
}

Verified that the adjacency matrix (adj_mx) is consistent in shape and values between MXNet and PyTorch. Ensured that the model architecture matches, including layers, activations (GLU), and the number of parameters. Used the same loss function (Huber Loss) and optimizer (Adam with lr=0.001) in both implementations. Checked the input data normalization to ensure consistency between frameworks. Confirmed that weight initialization in PyTorch uses Xavier initialization, as in MXNet.

What could cause such a significant performance gap between the PyTorch and MXNet implementations of the same model? Are there specific areas in PyTorch (e.g., matrix operations, batching, or gradient computation) that require attention when porting from MXNet? Any suggestions for debugging or aligning the two implementations more closely?

本文标签：