{"id":28411608,"url":"https://github.com/rishigurnani/polygnn_trainer","last_synced_at":"2025-06-23T23:30:57.450Z","repository":{"id":164468789,"uuid":"497720016","full_name":"rishigurnani/polygnn_trainer","owner":"rishigurnani","description":"A Python library to automate ML model training for polymer informatics.","archived":false,"fork":false,"pushed_at":"2024-02-15T23:33:55.000Z","size":208,"stargazers_count":5,"open_issues_count":2,"forks_count":4,"subscribers_count":1,"default_branch":"main","last_synced_at":"2025-06-03T05:27:51.969Z","etag":null,"topics":["data-science","machine-learning","materials-informatics"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/rishigurnani.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null}},"created_at":"2022-05-29T22:07:31.000Z","updated_at":"2025-03-15T04:18:54.000Z","dependencies_parsed_at":null,"dependency_job_id":"f906ec59-82dd-4d57-920a-ed6b432cfca3","html_url":"https://github.com/rishigurnani/polygnn_trainer","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/rishigurnani/polygnn_trainer","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/rishigurnani%2Fpolygnn_trainer","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/rishigurnani%2Fpolygnn_trainer/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/rishigurnani%2Fpolygnn_trainer/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/rishigurnani%2Fpolygnn_trainer/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/rishigurnani","download_url":"https://codeload.github.com/rishigurnani/polygnn_trainer/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/rishigurnani%2Fpolygnn_trainer/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":261575350,"owners_count":23179502,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["data-science","machine-learning","materials-informatics"],"created_at":"2025-06-02T17:42:52.781Z","updated_at":"2025-06-23T23:30:57.443Z","avatar_url":"https://github.com/rishigurnani.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# polygnn_trainer\nThis repository contains one of the custom packages used to train the machine learning models presented in a companion paper, [*polyGNN: Multitask graph neural networks for polymer informatics*](https://arxiv.org/abs/2209.13557). Currently, `polygnn_trainer` (pt) contains general code for the following tasks\n- Data preparation\n- Submodel training (and limited analysis of training metrics, provided by the [`parse`](https://github.com/rishigurnani/polygnn_trainer/tree/main/polygnn_trainer/parse) module)\n- Submodel saving \u0026 loading\n- Composing the submodels into an ensemble for inference\n## Installation\nThis repository is currently set up to run on 1) Mac OSX and 2) Linux/Windows machines with CUDA 10.2. Please raise a GitHub issue if you want to use this repo with a different configuration. Otherwise, please follow these steps for installation:\n\n1. Install [poetry](https://python-poetry.org/) on your machine.\n2. If Python3.7 is installed on your machine skip to step 3, if not you will need to install it. There are many ways to do this, one option is detailed below:\n    * Install [Homebrew](https://brew.sh/) on your machine.\n    * Run `brew install python@3.7`. Take note of the path to the python executable.\n3. Clone this repo on your machine.\n4. Open a terminal at the root directory of this repository.\n5. Run `poetry env use /path/to/python3.7/executable`. If you installed Python3.7 with Homebrew, the path may be something like\n  `/usr/local/Cellar/python\\@3.7/3.7.13_1/bin/python3.7`.\n7. Run `poetry install`.\n8. If your machine is a Mac, run `poetry run poe torch-osx`. If not, run `poetry run poe torch-linux_win-cuda102`.\n9. If your machine is a Mac, run `poetry run poe pyg-osx`. If not, run `poetry run poe pyg-linux_win-cuda102`.\n## Usage\n### General usage\nA usage pattern may look something like below:\n- First, import `polygnn_trainer`\n  ```python\n  import polygnn_trainer as pt\n  ```\n- Then, create a pandas Dataframe out of the data set (see [here](https://github.com/rishigurnani/polygnn_trainer/blob/main/sample_data/sample.csv) for an example CSV data set, note that you can distinguish between tasks using the \"prop\" column). Split this dataset into training+val and test set. It is required that graph features are provided in several dictionaries; one dictionary per row of the dataframe.\n- Pass the training+val DataFrame into `pt.train.prepare_train`. The outputs will be the dataframe, with SMILES converted to features, and dictionary of scalers. One scaler per task. Additionally, some metadata on training features will be saved to path input to `root_dir`. This will be useful to order features during inference.\n  ```python    \n  dataframe, scaler_dict = prepare_train(\n      \"dataframe\", morgan_featurizer, root_dir=/path/to/model/root/\n  )\n  ```\n- Split the training+val DataFrame into two DataFrames using your favorite splitting strategy (example shown below). One DataFrame for validation and one for training. See the example snippet below.\n  ```python    \n   training_df, val_df = sklearn.model_selection.train_test_split(\n      dataframe,\n      test_size=constants.VAL_FRAC,\n      stratify=dataframe.prop,\n      random_state=constants.RANDOM_SEED,        \n   )\n  ```\n- Run hyperparameter optimization using your favorite algorithm (example shown below). Use the error on the validation set to find the optimal set. See the example snippet below. The unittest named `test_ensemble_trainer` contains a more detailed, working, example. Unittests are located in [this folder](https://github.com/rishigurnani/polygnn_trainer/tree/main/tests).\n  ```python    \n  from skopt import gp_minimize\n  # obtain the optimal point in hp space\n  opt_obj = gp_minimize(\n      func=obj_func, # define this yourself\n      dimensions=hp_space,\n      n_calls=10,\n      random_state=0,\n  )\n  # create an HpConfig from the optimal point in hp space\n  hps = pt.hyperparameters.HpConfig()\n  hps.set_values(\n    {\n      \"r_learn\": 10 ** opt_obj.x[0],\n      \"batch_size\": opt_obj.x[1],\n      \"dropout_pct\": opt_obj.x[2],\n      \"capacity\": opt_obj.x[3],\n      \"activation\": torch.nn.functional.leaky_relu,\n    }\n  )\n  ```\n- Train an ensemble of submodels using `pt.train.train_kfold_ensemble`. Again, `test_ensemble_trainer` contains a nice working example. If the ensemble is being trained for research, we can train the ensemble using just the train+val set (so that we have the test set to assess model accuracy). If the ensemble is being trained for production, we can train the ensemble on the ***entire*** data set. The details of the model and its metadata will be stored in a root directory. The contents of that directory are as follows:\n   ```\n  root\n  │\n  └───metadata\n  │   │   features.pkl\n  │   │   features.txt\n  │   │   hyperparams.pkl\n  │   │   properties.txt\n  │   │   scalers.pkl\n  │   │   selectors.pkl\n  │   \n  └───models\n      │   model_0.pt\n      │   model_1.pt\n      │   ...\n  ```\n- Prepare a DataFrame of points to run inference on. Again, it is required that graph features are provided in several dictionaries; one dictionary per row of the melted dataframe.\n- Load the ensemble using `pt.load.load_ensemble`. Fields needed to instantiate the submodels can be passed into `submodel_kwargs_dict`.\n  ```python\n  ensemble = pt.load.load_ensemble(\n    \"path/to/model/root/\",\n    MyModelName,\n    device,\n    submodel_kwargs_dict={},\n  )\n  ```\n- Run inference on the ensemble using `pt.infer.eval_ensemble`. See below for example usage.\n  ```python\n  y_val, y_val_mean_hat, y_val_std_hat, selectors = pt.infer.eval_ensemble(\n    ensemble,\n    \"path/to/model/root/\",\n    dataframe,\n    smiles_featurizer,\n    device,\n    ensemble_kwargs_dict={\"monte_carlo\": False},\n  )\n  ```\n### `example.py`\nMuch of the information in the \"General usage\" section is combined into one file, `example.py`. Here we use training data located in the directory `sample_data` to train an ensemble model (composed of several submodels). The submodels, by default, are saved in a directory named `example_models`. The data in `sample_data` is a small subset of the DFT data used to train the models in the companion paper. A complete set of the DFT data can be found at [Khazana](https://khazana.gatech.edu/).\n\nTo train models run: `poetry run python example.py`. This should not take longer than 3 minutes on a machine with at least 8GB of free GPU memory. To manually specify the device you want to use for training, set the device flag. For example `poetry run python example.py --device cpu`. Otherwise, the device will automatically be chosen.\n\nLooking at `sample_data/sample.csv`, you will notice that this dataset contains multiple different properties (e.g., band gap, electron affinity, etc.). In `example.py`, we use this data to train a multitask model, capable of predicting each property. To train your own multitask model, you can replace `sample_data/sample.csv` with your own dataset containing multiple properties. Single task models are also supported.\n## Citation\nIf you use this repository in your work please consider citing us.\n```\n@article{Gurnani2023,\n   annote = {doi: 10.1021/acs.chemmater.2c02991},\n   author = {Gurnani, Rishi and Kuenneth, Christopher and Toland, Aubrey and Ramprasad, Rampi},\n   doi = {10.1021/acs.chemmater.2c02991},\n   issn = {0897-4756},\n   journal = {Chemistry of Materials},\n   month = {feb},\n   number = {4},\n   pages = {1560--1567},\n   publisher = {American Chemical Society},\n   title = {{Polymer Informatics at Scale with Multitask Graph Neural Networks}},\n   url = {https://doi.org/10.1021/acs.chemmater.2c02991},\n   volume = {35},\n   year = {2023}\n}\n```\n## License\nThis repository is protected under a General Public Use License Agreement, the details of which can be found in `GT Open Source General Use License.pdf`.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Frishigurnani%2Fpolygnn_trainer","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Frishigurnani%2Fpolygnn_trainer","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Frishigurnani%2Fpolygnn_trainer/lists"}