{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 训练模型\n",
    "\n",
    "基于 paddlex "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install paddlex -i https://mirror.baidu.com/pypi/simple\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://bj.bcebos.com/paddlex/datasets/optic_disc_seg.tar.gz\n",
    "!tar xzf optic_disc_seg.tar.gz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "matplotlib.use('Agg') \n",
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "import paddlex as pdx\n",
    "\n",
    "# 定义训练和验证时的transforms\n",
    "from paddlex.seg import transforms\n",
    "train_transforms = transforms.Compose([\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.Resize(target_size=512),\n",
    "    transforms.RandomPaddingCrop(crop_size=500),\n",
    "    transforms.Normalize()\n",
    "])\n",
    "eval_transforms = transforms.Compose([\n",
    "    transforms.Resize(512),\n",
    "    transforms.Normalize()\n",
    "])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = pdx.datasets.SegDataset(\n",
    "    data_dir='optic_disc_seg',\n",
    "    file_list='optic_disc_seg/train_list.txt',\n",
    "    label_list='optic_disc_seg/labels.txt',\n",
    "    transforms=train_transforms,\n",
    "    shuffle=True)\n",
    "eval_dataset = pdx.datasets.SegDataset(\n",
    "    data_dir='optic_disc_seg',\n",
    "    file_list='optic_disc_seg/val_list.txt',\n",
    "    label_list='optic_disc_seg/labels.txt',\n",
    "    transforms=eval_transforms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_classes = len(train_dataset.labels)\n",
    "num_classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_classes = len(train_dataset.labels)\n",
    "model = pdx.seg.FastSCNN(num_classes=num_classes)\n",
    "model.train(\n",
    "    num_epochs=40,\n",
    "    train_dataset=train_dataset,\n",
    "    train_batch_size=4,\n",
    "    eval_dataset=eval_dataset,\n",
    "    learning_rate=0.01,\n",
    "    save_interval_epochs=1,\n",
    "    save_dir='output/',\n",
    "    use_vdl=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!paddlex --export_inference --model_dir=./output/best_model --save_dir=./inference_model --fixed_input_shape=[512,512]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}