当前位置：首页 > news >正文

通过Dockerfile构建Docker镜像并训练模型

news 2025/9/13 7:33:00

1.准备工作

文件结构

models：存放预训练模型
src：存放运行脚本
requirements.txt：存放需要下载的环境
build.sh ：构建代码

Dockerfile实现

将文件、环境、代码等复制到docker容器中

FROM docker pull pytorch/pytorch:2.4.1-cuda12.1-cudnn9-runtime WORKDIR /app# 配置pip国内源
RUN pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple/ && \pip config set global.trusted-host pypi.tuna.tsinghua.edu.cn && \pip config set global.timeout 300# 升级pip
RUN pip install --no-cache-dir --upgrade pip setuptools wheel# 复制需求文件
COPY requirements.txt .# 安装Python依赖
RUN pip install --no-cache-dir -r requirements.txt || \pip install --no-cache-dir -r requirements.txt -i https://pypi.douban.com/simple/ || \pip install --no-cache-dir -r requirements.txt -i https://pypi.python.org/simple/# 复制应用代码 - 使用明确的路径
#COPY src/evaluate_model.py /app/
COPY ultralytics /app
COPY models /app
COPY src/train.py /app# 设置执行权限
RUN chmod +x /app/obb_train_finetune.py# 设置环境变量
ENV PYTHONPATH=/app
ENV PYTHONUNBUFFERED=1# 使用root用户运行（避免权限问题）
USER root# 设置入口点
ENTRYPOINT ["sh", "-c"]# 默认帮助信息
CMD ["--help"]

2.构建镜像

构建代码./build.sh

#!/bin/bashset -e# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color
APP_NAME="barcode-detection"echo -e "${GREEN} 开始构建${APP_NAME}应用...${NC}"# 检查 Docker 是否安装
if ! command -v docker &> /dev/null; thenecho -e "${RED} Docker 未安装，请先安装 Docker${NC}"exit 1
fi# 设置镜像名称和标签
IMAGE_NAME="ai-platform/${APP_NAME}"
TAG="latest"
FULL_IMAGE_NAME="${IMAGE_NAME}:${TAG}"echo -e "${YELLOW} 构建镜像: ${FULL_IMAGE_NAME}${NC}"# 构建 Docker 镜像
docker build --network=host -t "${FULL_IMAGE_NAME}" .if [ $? -eq 0 ]; thenecho -e "${GREEN} 镜像构建成功: ${FULL_IMAGE_NAME}${NC}"# 显示镜像信息echo -e "${YELLOW} 镜像信息:${NC}"docker images "${IMAGE_NAME}" --format "table {{.Repository}}\t{{.Tag}}\t{{.Size}}\t{{.CreatedAt}}"echo -e "${GREEN} 构建完成！${NC}"
elseecho -e "${RED} 镜像构建失败${NC}"exit 1
fi

3.训练和运行

构建

./build.sh

启动容器

docker run --rm -it --gpus all -v /mnt/nvme3n1/dujiang/obb_datasets/v8_refine_20241214/:/mnt/nvme3n1/dujiang/obb_datasets/v8_refine_20241214/  ai-platform/barcode-detection:latest bash -c