A knowledge distillation method and device of a generative adversarial network, equipment and storage medium

Abstract

The application discloses a knowledge distillation method of a generative adversarial network, and the method comprises the following steps: constructing a training data set, pre-processing an original image data set to obtain the training data set; training a teacher network, selecting a teacher network model according to an image generation task and training; constructing a student network, compressing a structure of an original generative adversarial network to obtain the student network and initializing the student network; distillation training of the student network, determining a knowledge distillation framework, training by adopting a short-term distillation strategy, and obtaining a light-weight generator. The training strategy of short-term distillation effectively improves the effect of the knowledge distillation training of the generative adversarial network, and further enhances the performance of the light-weight generator, so that a realistic high-resolution image can be generated. Moreover, the light-weight generator obtained by the application can run on a device with limited hardware resources, and has the advantages of low interaction delay and low memory occupation.

Description

TECHNICAL FIELD

[0001] The present application belongs to the technical field of computer vision, and particularly relates to a knowledge distillation method and device of a generative adversarial network, an equipment and a storage medium. BACKGROUND

[0002] In recent years, generative adversarial networks (GANs) are one of the most popular methods in the field of computer vision, and are widely used in image generation, style transfer, super-resolution, video synthesis and other tasks. With the increasing application of GANs, the demand for deploying generative adversarial networks on resource-constrained devices is gradually increasing. For example, improving video quality through super-resolution technology when the bandwidth is limited, generating images from users' stick figures, creating virtual face animations for remote social interaction, image deblurring, and drawing images from text. These tasks require users to interact with the model with low latency, but deploying network models on smart terminals such as mobile phones is restricted by hardware resources such as computing power and memory. Like other complex deep learning models, GANs have a huge parameter quantity and computational complexity, usually requiring a large storage space, a high computing cost and a large memory utilization, which makes them unable to be efficiently deployed, so it is necessary to obtain a lightweight generator through a specific method.

[0003] For common classification models, many mature network compression techniques have been developed to obtain efficient lightweight models, including pruning, quantization, knowledge distillation, neural network structure search, etc. However, compared with classification models, the structure and training process of GANs are more complex, and the compression difficulty is higher. Therefore, in order to reduce the inference time and model size of GANs while ensuring the generation effect, it is necessary to adjust and design a lightweight method specifically. Although some research work has achieved good results in the field of GANs lightweight, these works generally face conditional generation tasks at low resolution, such as pix2pix and CycleGAN. However, there are few works on high-resolution image generation with high realism. StyleGAN is one of the most popular generative adversarial networks, and has been widely used in academia and industry due to its excellent visual quality and excellent style decoupling ability. However, the calculation amount of StyleGAN for generating 1024*1024 resolution images reaches 144G MACs, which is more than 36 times of the commonly used classification network ResNet-50. Such a huge computing overhead seriously affects the deployment of StyleGAN on edge devices.

[0004] In order to solve the deployment problem of StyleGAN, a lightweight network is trained through knowledge distillation, the knowledge in the complex teacher network is transferred to the simple student network, and through the use of the output of the teacher network and even the intermediate layer information, the student network can more easily complete the training process. Due to the differences in structure and training, most of the existing distillation methods cannot be directly applied to StyleGAN, and the current method causes a large performance loss, therefore, according to the structure and characteristics of StyleGAN, the whole knowledge distillation process is redesigned, an efficient lightweight StyleGAN generator is obtained through a short-term distillation strategy, and the difficulties encountered in the application and deployment of StyleGAN are solved. SUMMARY

[0005] The application aims to provide a knowledge distillation method, device and equipment of a generative adversarial network and a storage medium, to solve the technical problems of large performance loss and huge computational overhead of StyleGAN caused by the existing distillation method.

[0006] The first object of the application is to provide a knowledge distillation method of a generative adversarial network, and the specific technical solutions of the application are as follows:

[0007] The knowledge distillation method of the generative adversarial network comprises the following steps:

[0008] Constructing a training data set, preprocessing the original image data set to obtain a training data set;

[0009] Training a teacher network, selecting a teacher network model according to an image generation task and training;

[0010] Constructing a student network, compressing the structure of the original generative adversarial network to obtain a student network and initializing the student network;

[0011] Distillation training of the student network, determining a knowledge distillation framework, adopting a short-term distillation strategy for training, and obtaining a lightweight generator.

[0012] Further, the specific steps of constructing the student network comprise:

[0013] Determining the compression mode of the student network, compressing the generator according to the characteristics of the generative adversarial network, and realizing it by reducing the convolution channel number of the generation network;

[0014] Determining the structure of the student network, training a generative adversarial network with variable channel number to reduce the convolution channel number according to the computational requirement, using the Inception Score to evaluate the generation quality of the network under different channel number configurations, searching for the optimal generator structure under the given computational requirement, establishing the generator of the student network according to the searched structure, and adopting the same setting as the teacher network for the discriminator structure;

[0015] The student network is initialized, and except that the mapping network in the generator needs to inherit the parameters in the teacher network, other parts are randomly initialized.

[0016] Further, the determination knowledge distillation framework includes that the student network learns the output of the teacher network through knowledge distillation while performing the adversarial training, requires the student network to obtain similar output to the teacher network under the same noise input, takes the pixel loss and the perceptual loss as the distillation loss to measure the output gap, and the student network is trained under the supervision of the teacher network and the discriminator, and through comparison with the real image and the teacher generated image, the student network can be avoided to generate images with fuzzy color and contour and irregularity, and the objective function of the generative adversarial network knowledge distillation training can be represented as:

[0017]

[0018] wherein, is the loss function of the adversarial training, and the pixel loss and the perceptual loss respectively, pixel and λ percept are weight coefficients of the two kinds of distillation losses, s is the student discriminator, s and G t represent the teacher generator and the student generator respectively, x is an image in a training data set, and z is normal noise, is the difference between the images generated by the teacher generator and the student generator, that is, the absolute value of the difference of each pixel value; in order to strengthen the attention of the knowledge distillation to the high-frequency information such as contour and texture in the image, the pre-trained LPIPS (Learned Perceptual Image Patch Similarity) model is adopted to calculate the similarity between the outputs of the teacher and the student as the perceptual loss In the LPIPS model, the distance between different images is measured by means of the VGG network, and then a multi-layer fully connected network is adopted to evaluate the similarity between images according to the distance, so that the result conforming to human visual perception can be obtained.

[0019] Further, the short-term distillation strategy includes the following two stages:

[0020] The first stage: simultaneously performing knowledge distillation and adversarial training;

[0021] The second stage: gradually stopping knowledge distillation and finally only performing adversarial training;

[0022] Therefore, the objective function of the training is :

[0023]

[0024] wherein, lambda is a weight coefficient of distillation loss, n represents the round of training (4000 pictures are input as a round), n k is the corresponding round of entering the second stage of training, n per represents the round of weight reduction, and epsilon is the proportional constant of distillation weight reduction; since the mapping of the normal noise of the adversarial training learning to the real image distribution is not fixed, the content of the image generated by the student network includes color, contour, texture and the like, which will naturally produce errors with the image generated by the teacher network, but the distillation loss caused by the content error of the generated image does not affect the reality of the generated image, therefore, when the distillation loss cannot be reduced with the training, it is considered that the divergence between the knowledge distillation and the adversarial training has begun to occur, at this time, the model is guided into the second stage of training.

[0025] The second object of the present application is to provide a training device for knowledge distillation of a generative adversarial network, the training device comprising:

[0026] a teacher network training module, configured to preprocess and data augment training data to obtain a training data set, and train an original StyleGAN network as a teacher network;

[0027] a student network construction module, configured to generate a corresponding student network structure according to a specified calculation amount requirement and initialize the student network structure;

[0028] a knowledge distillation training module, configured to establish a knowledge distillation framework, and perform short-term distillation training on the student network on the data set to obtain a deployable lightweight generator.

[0029] The third object of the present application is to provide a computer device, comprising a memory, a processor and a computer program stored in the memory and executable on the processor, wherein the processor implements the knowledge distillation method of the generative adversarial network when executing the computer program.

[0030] The fourth object of the present application is to provide a computer readable storage medium, the computer readable storage medium stores a computer program, and the computer program is executed by a processor to implement the knowledge distillation method of the generative adversarial network.

[0031] Compared with the prior art, the present application has the following advantages:

[0032] The student network construction method provided by the application can obtain a network structure with optimal performance under a given calculation requirement, reduces the calculation amount of the student generator, and ensures that the performance of the student generator is not damaged as much as possible; the short-term distillation training strategy improves the effect of the generative adversarial network knowledge distillation training, thereby enhancing the performance of the lightweight generator, and the lightweight generator can generate realistic high-resolution images; and the lightweight generator obtained by the application can run on a device with limited hardware resources, and has the advantages of low interaction delay and low memory occupation. BRIEF DESCRIPTION OF DRAWINGS

[0033] Figure 1 is a flowchart of the implementation of the knowledge distillation method;

[0034] Figure 2 is a schematic diagram of the generative adversarial network StyleGAN network adopted by the application;

[0035] Figure 3 is a schematic diagram of the knowledge distillation framework adopted by the application;

[0036] Figure 4 is a schematic diagram of the channel configuration search method adopted by the application;

[0037] Figure 5 is a schematic diagram of the composition structure of the generative adversarial network knowledge distillation training device. DETAILED DESCRIPTION

[0038] In order to make the technical means, creative features, purposes and effects of the application easy to understand, the application of a generative adversarial network knowledge distillation method, device, equipment and storage medium is further described in detail below in combination with specific drawings.

[0039] Embodiment one

[0040] As shown in Figure 1 , a generative adversarial network knowledge distillation method specifically includes:

[0041] S1: Construct a training data set, and pre-process an original image data set to obtain a training data set; specifically, the pre-processing of the original image data set includes cropping, aligning and data augmentation of the original image to obtain a data set that can be used to train the generative adversarial network.

[0042] S2: Train the teacher network, select a teacher network model according to the image generation task and train, and train the selected generative adversarial network StyleGAN on the data set to obtain a pre-training model as the teacher network used in this example.

[0043] S3: Constructing the student network according to the calculation requirement, obtaining the student network by reducing the convolution channel number of the original network generator, training the generative adversarial network with variable generator channel number as the search space of the student generator to obtain the generator structure meeting the calculation requirement, constructing the student network according to the channel structure and initializing the student network.

[0044] Specifically, the student network generator is obtained by reducing the channel number of the convolution network in the student network generator, a StyleGAN generator with variable channel number of each layer of convolution network is trained as the search space of the student network generator structure, and the channel compression rate of each layer of convolution network in the StyleGAN generator is set to 0.4, 0.5, 0.6, 0.7 and 1 according to experience. In the training process, the channel compression rate of each layer of convolution network is randomly selected in each round, and the convolution channels are sorted according to the weight size. According to the layer compression rate, the convolution channel with the smallest weight is excluded, so that the generator can complete the image generation task under different convolution channel configurations. In this example, the StyleGAN generator with variable channel number is obtained, and neural network search is performed using the StyleGAN generator to obtain the optimal channel configuration scheme under the specified calculation requirement.

[0045] As shown in Figure 2 , the channel structure of the generation network is searched by using an evolutionary algorithm in this example, InceptionScore is used as the evaluation function of the generation effect of each scheme, different channel setting schemes are used as individuals in the population, and a certain number of individuals are selected from the optimal part of the population in each round to occur hybridization and a certain number of individuals to occur mutation, and then the channel setting scheme meeting the calculation requirement is selected and added to the population. After a certain number of rounds, the best convolution network channel structure can be found.

[0046] As shown in Figure 3 , the generator of this example uses the searched structure, and the discriminator still uses the structure of StyleGAN. Based on the structure of StyleGAN, the channels of each layer of convolution network of the student generator are set to 512, 512, 512, 512, 358, 256, 256, 256, 128, 128, 64, 64, 32 and 32 according to the best structure under the calculation amount of 7.7 GMACs. The student network constructed in this example is randomly initialized except the mapping network, and the mapping network needs to inherit the parameters in the teacher network to ensure the smooth start of knowledge distillation.

[0047] S4: Distillation training of the student network, determining the knowledge distillation framework, using the short-term distillation strategy for training to obtain the lightweight generator.

[0048] As shown in Figure 4As shown, the student network is trained under the supervision of the teacher network and the discriminator, and the objective function of the generative adversarial network knowledge distillation training can be represented as:

[0049]

[0050] wherein, is the loss function of the adversarial training, and are the pixel loss and the perceptual loss respectively, λ pixel and λ percept are the weight coefficients of the two distillation losses, both of which are set to 3 in this embodiment, s is the student discriminator, G s and G t represent the teacher generator and the student generator respectively, x is an image in the training data set, and z is normal noise, is the error between the images generated by the teacher generator and the student generator, i.e., the absolute value of the difference of each pixel value; in order to strengthen the attention of the knowledge distillation to the high-frequency information such as contours and textures in the image, the present application adopts a pre-trained LPIPS (Learned Perceptual Image Patch Similarity) model to calculate the similarity between the outputs of the teacher and the student as the perceptual loss In the LPIPS model, the distance between different images is measured by means of a VGG network, and then a multi-layer fully connected network is used to evaluate the similarity between images according to the distance, so that the result consistent with human visual perception can be obtained.

[0051] The strategy of short-term distillation divides the training process into two stages according to the compatibility of knowledge distillation and adversarial training:

[0052]

[0053] In this example, the weight coefficient λ of the distillation loss is set to 3, the number of times n per is set to 25, the number of times n k corresponding to the second stage training is set to 2500, and the proportion constant ∈ of the distillation weight reduction is set to 0.9. The initial learning rate of the Adam optimizer is set to 0.002, and the two decay factors are set to 0 and 0.99 respectively.

[0054] Compared with the original distillation training method, the short-term distillation strategy is to only perform distillation at the initial stage of the lightweight generator training, and gradually exclude the influence of knowledge distillation when the model itself is in the right track of adversarial training, so the training process is divided into two stages, the first stage student network is trained under the joint action of knowledge distillation and adversarial training, and after entering the second stage, the distillation training is gradually stopped, in order to avoid the sudden stop of distillation leading to unstable training, the weight coefficient of distillation loss is proportionally reduced to 0 according to a certain number of rounds. The key to the training stage switching stage is the relationship between knowledge distillation and adversarial training, the compatibility of knowledge distillation and adversarial training is judged by the change rate of distillation loss, when the distillation loss cannot decrease with the training, it is considered that knowledge distillation and adversarial training have begun to have great divergence, in this example, the model is guided to enter the second stage of training after 2500 rounds of training.

[0055] The performance of the lightweight StyleGAN generator is tested by the index FID (Frechet Inception Distance) for evaluating the quality of the generated image of the generated model, and the lower the FID represents the more real the generated image of the generator, that is, the stronger the performance of the generator. As shown in Table 1, under the same network structure and data, the knowledge distillation method provided by the application obtains a lower FID, which significantly improves the performance of the generator. The calculation method of FID is as follows:

[0056] Randomly select 50000 real images to form a real image dataset r, and then use the generated image to generate 50000 images to form a generated image dataset g, input each image into the inception v3 model pre-trained on ImageNet, extract the activation value of the last layer of the model as the feature of the image, and calculate the mean value μ r 、μ g and variance Σ r 、Σ g of the image features in the two datasets respectively, and then calculate the Frechet distance between the images in the two datasets by the following formula, that is, the FID value (d) of the generator:

[0057]

[0058] Table 1 is the FID value comparison of the benchmark model and the knowledge distillation method of the application:

[0059] Experimental setup FID Baseline model 7.5 Baseline model + knowledge distillation method of the present application 5.4

[0060] Example two:

[0061] A training device for generative adversarial network knowledge distillation, the training device comprises:

[0062] The teacher network training module is configured to preprocess and augment data to obtain a training dataset, and train an original StyleGAN network as a teacher network.

[0063] The student network construction module is configured to generate a corresponding student network structure according to a specified computational requirement and initialize the student network structure.

[0064] The knowledge distillation training module is configured to establish a knowledge distillation framework, and perform short-term distillation training on the student network based on the dataset to obtain a deployable lightweight generator.

[0065] Embodiment three:

[0066] As shown in Figure 5 A computer device includes a memory, a processor, and a computer program stored in the memory and executable on the processor, and the processor implements the knowledge distillation method of the generative adversarial network in embodiment one when executing the computer program.

[0067] The memory is a volatile and / or non-volatile memory, and the memory can store data to support the operation of the teacher network training module, the student network construction module, and the student network training module; the processor reads executable instructions included in the software module in the memory, and combines necessary hardware to complete the knowledge distillation method of the generative adversarial network provided by the embodiments of the present application; after the student network training is completed, the generator part can be deployed on a device with limited computing resources.

[0068] Embodiment four:

[0069] A computer readable storage medium stores a computer program, and the computer program is executed by a processor to implement the knowledge distillation method of the generative adversarial network in embodiment one.

[0070] Finally, it should be noted that: the above embodiments are only used to illustrate the technical solutions of the present application, and not to limit them; although the present application has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that: it can still modify the technical solutions recorded in the foregoing embodiments, or make equivalent replacement to part or all of the technical features; and these modifications or replacements do not make the essence of the corresponding technical solutions deviate from the scope of the technical solutions of the embodiments of the present application.

Claims

1. A knowledge distillation method for generative adversarial networks, characterized in that, The method includes: Step 1: Construct the training dataset by preprocessing the original image dataset. Step 2: Train the teacher network. Select and train the teacher network model based on the image generation task. Step 3: Construct the student network by compressing the original generative adversarial network structure to obtain the student network and initializing it; the specific steps for constructing the student network include: The compression method for the student network is determined by the characteristics of generative adversarial networks, which dictate that the generator should be compressed by reducing the number of convolutional channels in the generator network. The structure of the student network is determined. In order to reduce the number of convolution channels according to the computational requirements, a generative adversarial network with a variable number of channels is trained. The Inception Score is used to evaluate the generation quality of the network under different channel configurations. The optimal generator structure under the given computational requirements is obtained by searching through an evolutionary algorithm. The generator of the student network is built based on the structure obtained by the search. The discriminator structure adopts the same settings as the teacher network. The student network is initialized, except for the mapping network in its generator, which needs to inherit the parameters from the teacher network; all other parts are randomly initialized. Step 4: Distillation training of the student network. Determine the knowledge distillation framework and use a short-term distillation strategy for training to obtain a lightweight generator. The knowledge distillation framework involves a student network learning the teacher network's output through knowledge distillation while undergoing adversarial training. The student network is required to produce outputs similar to the teacher network under the same noisy input. Pixel loss and perceptual loss are used as distillation losses to measure the output difference. The student network is trained under the supervision of the teacher network and a discriminator. The teacher network guides the training with real-image fusion, which effectively avoids the student network generating images with blurred colors and contours. The objective function for knowledge distillation training of the generative adversarial network is expressed as: in, It is the loss function for adversarial training. and These are pixel loss and perceptual loss, respectively, λ pixel and λ percept D is the weighting coefficient for these two types of distillation losses. s It is a student discriminator, G s and G t These represent the teacher generator and the student generator, respectively. x is an image from the training dataset, and z is normal noise. It is the error between the images generated by the teacher generator and the student generator, that is, the absolute value of the difference between each pixel value; The pre-trained LPIPS model is used to calculate the similarity between teacher and student outputs as the perceptual loss. The LPIPS model can extract image contour and texture information, thereby enhancing the effect of knowledge distillation. The short-term distillation strategy includes the following two stages: Phase 1: Simultaneous knowledge distillation and adversarial training; Phase Two: Gradually stop knowledge distillation and eventually only conduct adversarial training; Therefore, the objective function of training for: where λ is a weight coefficient of distillation loss, n represents the corresponding round of training, n k to enter the second phase of training corresponding round, n per represent the round of weight reduction, ∈ is the proportion constant of distillation weight reduction.

2. A training apparatus for knowledge distillation in generative adversarial networks, used to execute the knowledge distillation method for generative adversarial networks as described in claim 1, characterized in that, The training device includes: The teacher network training module is used to preprocess and augment the training data to obtain the training dataset, and to train the original StyleGAN network as the teacher network. The student network construction module is used to generate and initialize the corresponding student network structure according to the specified computational requirements. The knowledge distillation training module is used to establish a knowledge distillation framework and perform short-term distillation training on the student network on the dataset to obtain a deployable lightweight generator.

3. A computer device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, characterized in that, When the processor executes the computer program, it implements a knowledge distillation method for generative adversarial networks as described in claim 1.

4. A computer-readable storage medium storing a computer program, characterized in that, When the computer program is executed by the processor, it implements the knowledge distillation method for generative adversarial networks as described in claim 1.

Images