docker容器中训练tf模型

  • 2020-06-13
  • 79
  • 0

在 docker 下配置 cuda、tensorflow,以及训练模型(优点是多用户之间保持隔离)。docker 的安装见电子书 Docker – 从入门到实践/安装docker

1. 创建用户

sudo useradd -m potoo  # 创建 potoo,
sudo usermod -aG sudo potoo  # 并分配到sudo组,使其可用sudo提权
sudo passwd potoo  # 更改 potoo 的密码

su potoo  # 切换到 potoo 用户
cd ~  # 切换到当前用户的家目录

2. 安装 cuda

使用 Nvidia 容器工具包 来配置环境。此处从其官网中截取 nvidia-docker 的安装命令:

# 添加 Nvidia 容器工具源地址
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list

# 安装 Nvidia 容器工具包
sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit

# 重启 docker
sudo systemctl restart docker

docker 拉取 cuda 镜像:

docker pull nvidia/cuda:latest

测试 cuda:

docker run --gpus all --rm nvidia/cuda:latest nvidia-smi

注意:截止 docker == 19.03.8, docker-compose == 1.25.5,还未正式支持 docker-compose 部署和启动 nvidia-docker,相关 issue: Support for NVIDIA GPUs under Docker Compose。有个曲折的方法是:通过 nvidia-container-runtime 设置启动前的 hook。

3. 安装 Tensorflow-GPU

Tensorflow 安装详细见官方文档: Tensorflow Install/docker

TensorFlow Docker images 镜像拉取,由于我下面要使用 tfhub V2,故这里拉取最新镜像(py 3.6, tf 2.x) (自从 220.1.27 起有标签 py3 和无标签的均指向同一份 py3 的镜像):

docker pull tensorflow/tensorflow:latest-gpu-jupyter

关于 tag 带 jupyter 的镜像,查看其镜像构建会发现其会自动启动 jupyter-notebook,并且目录为 /tf,镜像构建的最后一条命令如下:

/bin/bash -c #(nop)  CMD ["bash" "-c" "source /etc/bash.bashrc && jupyter notebook --notebook-dir=/tf --ip 0.0.0.0 --no-browser --allow-root"]

测试:

# 运行镜像并进入容器内
docker run --gpus all -it tensorflow/tensorflow:latest-gpu-jupyter /bin/bash

# 查看 python 与 tensorflow 版本
python -V  #  本次版本为 Python 3.6.9
pip list | grep tensorflow  # 本次版本为 2.2.0
jupyter --version  # jupyter notebook 为 6.0.3

# 测试 GPU,能正确看到 GPU 信息说明以上环境全部成功配置
python -c "import tensorflow as tf; print(tf.reduce_sum(tf.random.normal([1000, 1000])))"

4. 训练

先介绍下我所知的开源图像分类模型训练方法,有以下两种方法:

我对这两种使用开源模型的优缺点看法:

  • hub:
    • 优点:训练十分简单,数据集不需要额外处理;
    • 缺点:版本太高
  • slim:
    • 优点:开源模型很丰富;
    • 缺点:图像数据需要转化成 tfrecord。

4.1 文件准备

此处使用 hub V2,下载官方提供 notebook 例程 以及模型,国内无法访问 tfhub,所以此处先下载模型解压使用。本次使用模型为 imagenet 预训练的 mobilenet V2,下载到 models 目录下并解压:

tar zxvf ./models/imagenet_mobilenet_v2_140_224_classification_4.tar.gz -C ./models

启动 #3 拉取的镜像,将 V2 的 notebook 例程、下载的模型以及自己的数据集映射到容器的 /tf 文件夹下即可,如:

docker run \
    --name trainTfModel \
    -p 8888:8888 \
    -v /home/potoo/trainTfModel:/tf/trainTfModel \
    --gpus all \
    tensorflow/tensorflow:latest-gpu-jupyter

其中我的 trainTfModel 目录结构:

$ tree trainTfModel -L 2 -p
trainTfModel
├── [drwxr-xr-x]  datasets
│   ├── [drwxr-xr-x]  fish
│   ├── [drwxr-xr-x]  garbage
│   └── [drwxr-xr-x]  ship
├── [drwxr-xr-t]  models
│   ├── [drwxr-xr-t]  assets
│   ├── [-rwxr-xr-x]  imagenet_mobilenet_v2_140_224_classification_4.tar.gz
│   ├── [-rw-r--r--]  saved_model.pb
│   └── [drwxr-xr-t]  variables
├── [drwxr-xr-x]  output
│   └── [drwxr-xr-x]  saved_model
├── [-rw-r--r--]  tf2_image_fineTune.ipynb
└── [-rw-r--r--]  tf2_image_retraining-Copy1.ipynb

9 directories, 4 files

在官方提供的 notebook 例子中需要修改自己数据集路径和模型路径。

4.2 notebook 中训练

训练之前安装 tfhub, pillow,此处均在 notebook 中进行:

!pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -U tensorflow_hub pillow

以及我这里出现的错误: OSError: image file is truncated,猜测是因为服务器允许的单次数据大小问题,详细讨论见 Python PIL “IOError: image file truncated” with big images,此处引用其解决方案:

from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True

更改模型路径以及输入图像尺寸,根据上面我的 trainTfModel 目录结构,此处更改数据集路径为:

MODULE_HANDLE = "./models"
IMAGE_SIZE = (224, 224)

更改自己数据集路径,根据上面我的 trainTfModel 目录结构,此处更改数据集路径为:

data_dir = './datasets'

最后按照自己需要修改是否 Data Augmentation, Fine-Tune。

我的 notebook 完整如下:

# Retraining an Image Classifier

## 1. 环境配置

安装镜像中未安装的 tensorflow_hub 与 pillow,以及我这里出现的错误: *OSError: image file is truncated*,猜测是因为服务器允许的单次数据大小问题,详细讨论见 [Python PIL “IOError: image file truncated” with big images](https://stackoverflow.com/questions/12984426/python-pil-ioerror-image-file-truncated-with-big-images)。

​```python
!pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -U tensorflow_hub pillow
​```

​```python
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
​```

​```python
import itertools
import os

import matplotlib.pylab as plt
import numpy as np

import tensorflow as tf
import tensorflow_hub as hub

print("TF version:", tf.__version__)
print("Hub version:", hub.__version__)
print("GPU is", "available" if tf.test.is_gpu_available() else "NOT AVAILABLE")
​```

## 2. 模型加载
这里是自己手动下载的方式。下载后需要解压。取消下面一行的注释并执行即可。

​```python
#!tar zxvf ./models/imagenet_mobilenet_v2_140_224_classification_4.tar.gz -C ./models
​```

​```python
!ls models
​```

​```python
MODULE_HANDLE = "./models"
IMAGE_SIZE = (224, 224)
print("Using {} with input size {}".format(MODULE_HANDLE, IMAGE_SIZE))

BATCH_SIZE = 16 #@param {type:"integer"}
​```

## 3. 数据集处理

将图片缩放裁剪以符合模型输入,以及划分训练验证集、打乱顺序、数据集增强。

​```python
data_dir = './datasets'
#!ls ./datasets
​```

​```python
datagen_kwargs = dict(rescale=1./255, validation_split=.20)
dataflow_kwargs = dict(target_size=IMAGE_SIZE, batch_size=BATCH_SIZE,
                   interpolation="bilinear")

valid_datagen = tf.keras.preprocessing.image.ImageDataGenerator(
    **datagen_kwargs)
valid_generator = valid_datagen.flow_from_directory(
    data_dir, subset="validation", shuffle=False, **dataflow_kwargs)

do_data_augmentation = True #@param {type:"boolean"}
if do_data_augmentation:
    train_datagen = tf.keras.preprocessing.image.ImageDataGenerator(
        rotation_range=40,
        horizontal_flip=True,
        width_shift_range=0.2, height_shift_range=0.2,
        shear_range=0.2, zoom_range=0.2,
        **datagen_kwargs
    )
else:
    train_datagen = valid_datagen

train_generator = train_datagen.flow_from_directory(
    data_dir, subset="training", shuffle=True, **dataflow_kwargs)
​```

## 4. 模型定义

丢弃原模型的输出层,新建输出层以适应自己数据集的种类。

​```python
do_fine_tuning = True #@param {type:"boolean"}
​```

​```python
print("Building model with", MODULE_HANDLE)
model = tf.keras.Sequential([
    # Explicitly define the input shape so the model can be properly
    # loaded by the TFLiteConverter
    tf.keras.layers.InputLayer(input_shape=IMAGE_SIZE + (3,)),
    hub.KerasLayer(MODULE_HANDLE, trainable=do_fine_tuning),
    tf.keras.layers.Dropout(rate=0.2),
    tf.keras.layers.Dense(train_generator.num_classes,
                          kernel_regularizer=tf.keras.regularizers.l2(0.0001))
])
model.build((None,)+IMAGE_SIZE+(3,))
model.summary()
​```

## 5. 模型训练

### 5.1 优化器

此处使用默认的 SGB,随机梯度下降法。

​```python
model.compile(
    optimizer=tf.keras.optimizers.SGD(lr=0.005, momentum=0.9), 
    loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True, label_smoothing=0.1),
    metrics=['accuracy']
)
​```

### 5.2 训练

​```python
steps_per_epoch = train_generator.samples // train_generator.batch_size
validation_steps = valid_generator.samples // valid_generator.batch_size
hist = model.fit(
    train_generator,
    epochs=5, steps_per_epoch=steps_per_epoch,
    validation_data=valid_generator,
    validation_steps=validation_steps).history
​```

### 5.3 loss、accuracy 曲线

​```python
plt.figure()
plt.ylabel("Loss (training and validation)")
plt.xlabel("Training Steps")
plt.ylim([0,2])
plt.plot(hist["loss"], label='loss')
plt.plot(hist["val_loss"], label='val_loss')
plt.legend();

plt.figure()
plt.ylabel("Accuracy (training and validation)")
plt.xlabel("Training Steps")
plt.ylim([0,1])
plt.plot(hist["accuracy"], label='accuracy')
plt.plot(hist["val_accuracy"], label='val_accuracy')
plt.legend();
​```

### 5.4 保存模型

​```python
saved_model_path = "./output/saved_model"
tf.saved_model.save(model, saved_model_path)
​```


版权声明: 本网站所有资源采用BY-NC-SA 4.0协议进行授权,转载应当以相同方式注明文章来自:docker容器中训练tf模型 - 一方的天地

评论

还没有任何评论,你来说两句吧

发表评论

陕ICP备18010914号
知识共享许可协议
本作品由一方天地采用知识共享署名-非商业性使用-相同方式共享 4.0 国际许可协议进行许可,转载或引用本站文章应遵循相同协议。如果有侵犯版权的资源请尽快联系站长,本站会在24h内删除有争议的资源。 -