Detailed microstructure evolution in shape memory alloys (SMAs) is typically studied by molecular dynamics (MD) simulations. However, the conventional post-processing tools for atomistic calculations, such as CNA and PTM, fail to identify distinct crystal variants and to reveal twin alignments in SMAs. In the current work, a powerful and efficient post-processing tool based on GraphSAGE neural network is developed, which can identify multiple phases in martensitic transformation, including the orthorhombic, monoclinic and R phases. Where the network was trained by the results of sets of temperatureand stress-induced martensitic transformation MD calculations. The accuracy and generality were also verified by the application to the cases which did not appear in the training dataset, such as unrecoverable nanoindentation process. The proposed method is rapid, accurate, and is ready to be integrated with any visualization tool for MD simulations. The outcome of the current work is expected to accelerate the pace of atomistic studies on SMAs and related materials.