Identification of proton and gamma plays an essential role in ultra-high energy gamma-ray astronomy with LHAASO-KM2A. In this work, two neural networks (deep neural networks (DNN) and graph neural networks (GNN)) are applied to distinguish proton and gamma in the LHAASO-KM2A simulation data. The receiver operating characteristic (ROC) curves are used to evaluate the quality of the model. Both KM2A-DNN and KM2A-GNN models give higher Area Under Curve (AUC) scores than the traditional baseline model.