UCloud AI Train Deep Learning Practice: Fast style transfer using TensorFlow

This article describes how to implement a fast style migration project and step by step online training via the local machine and UCloud AI Train service. At the end of the training, the input images were transferred and output into images. This article also compares the time/price of using a local machine versus UCloud AI Train.

A previous post described using Pytorch to achieve prisma-like effects, but the comments section was full of complaints that it was too slow to run on their own computers. On my machine at the time, an image could take up to five or six minutes to complete. This is because every time we input two pictures into the system, we must first process the style file and train the network that can effectively generate pictures. This step is independent of the target picture. If we can create this model, save it, and then every time a new image comes in, we don’t retrain the model, which can greatly reduce the execution time.

In fact, when PRISMA was initially given only a few styles to choose from, training ahead of time could significantly save the server money and reduce user wait times. There are also many subsequent articles and codes to realize early training, such as Perceptual Losses for real-time Style Transfer and super-resolution and the deep learning practice of Zhihu: realizing rapid Style Transfer using Tensorflow

This time, we will implement a faster application that can pre-train image style transfer in this way. The main body of this part of the code is zhihu @Hezhiyuan Daishen’s code, but the network was changed from VGG 16 to VGG 19, and some other networks were tested and compared.

Most machine learning problems can be divided into two parts: training and execution. This paper also explains these two parts respectively. In the training part, the style image file and the specified parameters are passed in, and the model file is generated by TensorFlow. In the execution part, the target image is passed in, and the corresponding transformed image is generated. Finally, a simple API is implemented for us to call.

training

The training of machine learning task is a process of extracting features from data, which often needs the support of a large number of data sets. Before choosing the data set, it is important to explain that the task of style migration is very similar to image recognition. There were two major breakthroughs in the initial style transfer. One was to use deep learning to obtain the texture of the image, while the previous work was all manual modeling. Textures in images, on the other hand, are described by local statistical features. For example:

We can think of this image as the texture of an apple, described geometrically as green or red, with a circular pattern and a dot in the middle. These can all be considered local features. In the object recognition task, a large number of local feature recognizers are used, such as VGG 19. So we can consider using object recognition models and data sets directly to extract local textures, or local features. More specifically, we used COCO Dataset, which is a Dataset that can be used for image classification and background segmentation.

With that in mind, we can start our training. The training process can be divided into the following steps:

1. Download the dataset file
2. Download the network model file
3. Write/download the training code
4. Train on the machine/use UCloud AI Train service to Train

First we download the compressed data set here (~13 GB). Unzip it and you get a train 2014 folder.

In addition, we also need to download the VGG 19 network model, which can be downloaded here. If you are interested in experimenting with other networks, you can also look for Models such as ResNet and Inception in Pretrained Models. This article uses VGG 19 and VGG 16 networks.

I have also put relevant data sets and model files on Baidu Cloud, which should be faster to download in China. You can download them here

I’ve already posted the training code on Github. Let’s clone it first:

git clone https://github.com/stevefermi/style-transfer.git
Copy the code

Create a new folder called Pretrained in this directory. Put our downloaded datasets train 2014 and vGG_19.ckpt (or vGG_16.ckpt, similar) into this folder. The resulting folder structure is as follows:

- style transfer
 - pretrained
     - train2014
     - vgg_19.ckpt
 - train.py
 - eval.py
 - conf
 - others...
Copy the code

Then execute Python train.py in the style Transfer level directory to begin training. Conf/Mosaic. Yml is selected by default as a style-style image. Py -c conf/candy. Yml if you want to train different styles of images, use -c to specify python train.py -c conf/candy.

If you want to visualize the training process in real time, TensorFlow provides a very useful component of the TensorBoard. Enter the models folder generated in the training process, use TensorBoard –logdir= Mosaic, and enter 127.0.0.1:6006 in the browser to view all Summary provided to us by TensorBoard, including various loss, as shown below:

If you’re using a powerful NVIDIA graphics card (GTX 1060 or above), you can turn on CUDA acceleration, which can significantly reduce computing time. With my GTX 1070+ 16GB of RAM + I7 6700 processor, each step takes about 0.6 seconds. It takes about 10 to 12 hours to get a good result. If you can’t wait, download the trained model file. Of course, with the development of cloud computing, training on a server becomes a cheaper and more cost-effective solution. Next, we tried to use the UAI-train service provided by UCloud. UAI provided the single card/four card scheme of P40, and the price was 0.17/0.6 yuan per minute. Compared with buying machines for training, the cost was much lower.

The training process of UAI on UCloud can be divided into the following steps:

1. Account registration and real name verification
2. File structure
3. Upload data sets to UFile
4. Package the code image and upload it
5. Creating training Tasks
6. Look at TensorBoard and download the trained model file

The first step is to register your account and verify your real name. First register a UCloud account. After registration and real-name verification, you can obtain Public and Private keys as shown below:

After getting it, we also need to create a Docker image to store our code and data. Select UHub from the menu above and create a Docker image.

The second step is to clean up the structure of the file, move the Pretrained folder out and rename it as data folder, and rename the style-Transfer folder as Code folder (it’s ok to leave the name unchanged, but keep the same directory structure).

The third step is to upload the data set to UFile. Firstly, we create a UFile Bucket. Since we upload a large number of files (~ 580,000 files), we need to use a file management tool to upload the files. Therefore, we need to select a file management tool suitable for our system to download the Bucket.

After downloading, we can modify the contents of config. CFG by changing public_key and secret_key to our own. And then we can run it

./filemgr-win64.exe --action mput --bucket demobucket --key key --dir train2014 --prefix train2014/ --trimpath <path to train2014 folder>  
Copy the code

If the result is as shown in the figure:

That means our data set starts to upload, which can take anywhere from 30 minutes to an hour, depending on the Internet speed. Adding a powerful VGG module that I’ve captured during this test is a powerful feature of the VGG module that I capture (CKPT), adding a powerful VGG module that I capture during this test

Step 4: Package the code image and upload it. This is a crucial step.

4.1 Packing and uploading requires the SDK of uAI-train service. After ensuring that Git has been installed, you can use the following command in a directory:

git clone https://github.com/ucloud/uai-sdk.git
cd uai-sdk
sudo python setup.py install # Decide if you need to use sudo
Copy the code

Copy the uaitrain_tool/tf/tf_tool.py file to the same folder as code. The final code and data structure is as follows:

- code
    - train.py
    - others...
- data
    - train2014
    - others...
- tf_tool.py
Copy the code

4.2 Starting to Pack. With tf_tool.py in place and the file structure sorted out, we can start packing. We can create a.bat(Windows) file or a.sh(Linux) script for packaging, which looks like this:

python tf_tool.py pack --public_key=<public_key> --private_key=<private_key> --code_path=./code/ --mainfile_path=train.py --uhub_username="<username>" --uhub_password="<password>" --uhub_registry=<your_image_name> --uhub_imagename=<your_image_name> --uhub_imagetag=<your_image_tag> --internal_uhub=false- ai_arch_v tensorflow - 1.3.0 -- test_data_path =. = / data -- test_output_path =. / data/models - train_params =""- python_version = = python - 3.6.2 - the OS ubuntu 16.04Copy the code

Test_data_path and test_output_path indicate the location of the data set and exported model file in the local test. Absolute paths are recommended. You need to replace the public/private key with the previously obtained private key and the uhub-Registry with the image name. Uhub_username and password are the UCloud account passwords. Image_tag is a tag file that can write information such as version numbers

UCloud’s packaging tool will mount the data folder we passed in under the /data/data folder of the container, so change the path in our training and configuration files accordingly. /data/data/train2014 pretrained/ trained/train2014 loss_model_file conf/ Mosaic. “Pretrained /vgg_19. CKPT” # the path to the checkpoint changed to loss_model_file: “/data/data/vgg_19.ckpt” # the path to the checkpoint

Also, remember to change the path to save the model to /data/output. In our code, training_path changed to training_path = os.path.join(‘/data/output’, flags.naming) in line 23

After that, we are ready to run our packaging tool. Run the previous script if it returns:

It proves that our packaging process is correct.

4.3 Viewing the Packing Result. After wrapping, we can see files such as uaitrain_cmd.txt are automatically generated in the directory. This file contains the commands for local testing and deployment, and after the run, the following files are generated:

CMD Used for deploying: /data/train.py 
CMD for CPU local test: sudo docker run -it -v D:\Learn\tensorflow\uai-prisma\data:/data/data -v D: \ Learn \ tensorflow \ uai - prisma \ data \ models: / data/output style_transfer - CPU: v0.1 / bin/bash - c"cd /data && /usr/bin/python /data/train.py --work_dir=/data --data_dir=/data/data --output_dir=/data/output --log_dir=/data/output"
CMD for GPU local test: sudo nvidia-docker run -it -v D:\Learn\tensorflow\uai-prisma\data:/data/data -v D: \ Learn \ tensorflow \ uai - prisma \ data \ models: / data/output uhub. Ucloud. Cn/style_transfer/style_transfer: v0.1 / bin/bash - c"cd /data && /usr/bin/python /data/train.py --num_gpus=1 --work_dir=/data --data_dir=/data/data --output_dir=/data/output --log_dir=/data/output"

Copy the code

Step 5: Create the AI training task. Go back to UCloud’s console and select AI Training Services. In creating the AI training task, because of the need to download the data in the Ufile, so you need to fill in the address and Public/Secret Key, and will be uaitrain_cmd. TXT CMD in this file, informs the for deploying content can be copied.

Step 6: Review the training process and results. After the training task is created, you can see the log of the training process in the background. The logs are refreshed every 5 seconds. If any problems are found, you can stop the task, modify it, and run it again.

As you can see from the log, uAI-train first downloads the files required for training, and then starts the training process. Once you’ve made sure it starts up properly, you can relax and take a coffee break before coming back 😀

During the execution of the task, the button of TensorBoard will appear on the right side of the task. Without any configuration, click this button to see the details of TensorBoard:

After running about 40000 steps, we finally got the pre-trained model we wanted. At this time, Loss was basically convergent. Upon completion, our training program will stop automatically and save the model in the Models folder. During my own training, TensorBoard can be used to check the result in real time. If the loss does not change much and is relatively stable, it can be stopped.

After the training, we can see the cost of the training process and see the corresponding log and model file in the UFile:

I slightly compared the time and cost of the UAI-train in my GTX 1070 in single-card mode for your reference.

The environment	Average time per step	The cost of	The total available
GTX 1070	0.6	About 3000 yuan per card	4 hours
UAI single card	0.47	About 30 yuan	3.2 hours

perform

Once you have the model, it’s much easier to execute. Find the eval.py file in your code and use the following command:

python eval.py --model_file .\models\mosaic\fast-style-model.ckpt-done --image_file .\img\test.jpg
Copy the code

Where image_file can be specified as another image, you can generate the image after the style migration. After execution we can look at the output time:

As you can see, it basically takes about 1 second to complete, which is a huge improvement compared to the previous one, which takes nearly 10 minutes for each sheet. Due to the increase in iteration steps, the effect is also very good:

If you are not satisfied with the result, you can also change the corresponding style_weight value in the conf/*.yml file. Reduce style_weight if there are too many textures to find the content of the image, or increase style_weight if there are too few textures. Find an optimal range by constantly tweaking and experimenting.

Something to watch out for

In the input and output positions of the file, be careful to read and write from the mount point. For example, the input path of the model is /data/output and /data/data/train2014. If not, the file cannot be read or the trained model cannot be uploaded to UFile.