OLP-GNN

This repository contains the code of our paper "Learning Optimal Linear Precoding for Cell-Free Massive MIMO with GNN", in Joint European Conference on Machine Learning and Knowledge Discovery in Databases (ECML PKDD), 2024. The datasets are available at this address: https://doi.org/10.6084/m9.figshare.25250890.

© 2024 Nokia
Licensed under the BSD 3-Clause Clear License
SPDX-License-Identifier: BSD-3-Clause-Clear

Dependencies

This program is implemented in Python 3.10 with PyTorch 2.0.0, Pytorch Geometric 2.3.0 and PyTorch Lightning 2.0.2. The full list of packages can be found in requirements.txt. Run pip install -r requirements.txt to install them.

PyPAPI (python_papi) is an optional package only required to count the number of FLOPs of the MR, ZF, OLP precoders in data_generation and OLP-GNN in plot_flops.py. It is included in requirements.txt.

In addition to these dependencies, the scripts in data_generation require to install MOSEK as the SOCP solver called by cvxpy. The instructions to install MOSEK can be found on their website: https://docs.mosek.com/latest/install/installation.html.

Quickstart

Here are quick instructions to reproduce the numerical results presented in our paper.

1. First download the raw datasets from https://doi.org/10.6084/m9.figshare.25250890 and place them in a separate folder, e.g., named data.

2. Run the following script to preprocess and split the data into training, validation and test datasets, which are saved in a single file dataset_train.pt for later use:

   python data_preprocessing.py data

The model proposed in the paper is saved as a Pytorch Lightning checkpoint file in trained_model/checkpoints/best_epoch=953.ckpt. The following scripts evaluate this model on the test data in terms of spectral efficiency, number of FLOPs and runtime. In the next steps 3-5, we save all numerical results and figures in a folder figures.

3. The spectral efficiency CDFs for all test datasets are obtained by:

   python plot_cdf.py trained_model/checkpoints/best_epoch=953.ckpt data figures

4. The FLOPs are counted by the following script. Running this script requires python_papi installed and a CPU compatible with PAPI_SP_OPS event.

   python plot_flops.py trained_model/checkpoints/best_epoch=953.ckpt data figures

5. The runtimes are measured in:

   python test_compiled_runtimes.py trained_model/checkpoints/best_epoch=953.ckpt figures

Repository Overview

Data Generation

Folder data_generation contains the code to simulate various LoS and NLoS channel coefficients and generate MR, ZF and OLP precoding matrices. The data used in the paper is available in https://doi.org/10.6084/m9.figshare.25250890. Additional data could be produced with the following calls:

LoS data can be generated by running: python data_generation_los.py <N> <M> <K> <f> <precoder>.

NLoS data can be generated by running: python data_generation.py <N> <M> <K> <mor> <precoder>.

N is the number of graphs to generate, M the number of APs, K the number of UEs. In the LoS case, f is the carrier frequency in GHz. In the NLoS case, mor is the radio propagation morphology which can be either urban, suburban or rural. In both cases, precoder refers to the precoding algorithm, which can be either olp, mr or zf.

Graph Neural Network

There are 3 GNN implementations in this repository: GNNLinearPrecodingPyG is a Lightning Module implemented with Pytorch Geometric functions. FastGNNLinearPrecoding is a pure Pytorch module (does not use Lightning or PyG), it can be compiled and has the fastest inference runtime among all. FastGNNLinearPrecodingLightning is a Lightning model that encapsulates FastGNNLinearPrecoding for simpler developpement routines (training, validation and testing with logs and checkpoints).

FastGNNLinearPrecoding and GNNLinearPrecodingPyG are similar but their outputs may differ due to different implementations of the basic operations (PyTorch vs. PyTorch Geometric).

Training Pipeline

Before training the model, one must first preprocess the raw data files generated in the data_generation folder. The script data_preprocessing.py preprocesses and splits the data into training, validation and test datasets. Any new raw data file must be added to this script with its split (train_split, validation_split, test_split) specified.

python data_preprocessing.py <raw_data_folder>

Then, to train an OLP-GNN model:

python gnn_train.py

Finally, the obtained model checkpoint (.ckpt) can be evaluated following the instructions in the above quickstart section.