bitspeech/tensor2tensor

镜像概述

本镜像为Tensor2Tensor（TensorFlow生态下的机器学习工具）的容器化版本，旨在简化机器学习模型的训练与生成流程，支持CPU和GPU（CUDA）两种运行环境。

核心功能

• 提供CPU与GPU（CUDA）版本镜像，适配不同硬件环境
• 集成t2t-trainer工具，支持模型训练、数据生成与评估
• 快速启动容器并暴露端口，方便进行交互式操作

使用场景

• 机器学习研究者快速原型开发与模型测试
• 开发者进行计算机视觉（如MNIST图像识别）、自然语言处理等任务的模型训练
• 教学或学习中演示机器学习模型的训练流程

配置说明

• CPU版本：直接使用常规Docker命令启动
• GPU版本：需先安装nvidia-docker工具，以支持CUDA加速

部署示例

CPU容器启动

bash
docker run -it -p 8888:8888 bitspeech/tensor2tensor:1.6.6 /bin/bash

GPU（CUDA）容器启动

先安装nvidia-docker，再执行以下命令：

bash
nvidia-docker run -it -p 8888:8888 bitspeech/tensor2tensor:1.6.6-gpu /bin/bash

快速开始

使用t2t-trainer进行MNIST模型训练示例：

bash
t2t-trainer \
  --generate_data \
  --data_dir=~/t2t_data \
  --output_dir=~/t2t_train/mnist \
  --problem=image_mnist \
  --model=shake_shake \
  --hparams_set=shake_shake_quick \
  --train_steps=1000 \
  --eval_steps=100

更多详情请参考官方文档：https://github.com/tensorflow/tensor2tensor