Computer vision technology is widely used in autopilot system, which mainly solves the problem of object recognition and object classification. In this study, a lightweight neural network structure is proposed according to the task. In order to solve the problem of insufficient training data, an improved data enhancement algorithm is adopted to double the training data. At the same time, in order to solve the problem of using data generator as verification set and unable to use tensorboard, a solution is proposed. The principle of neural network processing image information is studied in detail by convolution network visualization method, and the optimization method is put forward. The accuracy of the trained model is 97.5% on the verification set, which meets the accuracy needs of autopilot system for classification task.