{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# 简介","metadata":{}},{"cell_type":"markdown","source":"* 使用 pytorch 自带的 fastrer-rcnn， 训练2小时，能拿到公榜 mAP 58.9 的效果\n* notebool 展示了使用 torchvision 中 fastrer-rcnn，验证集取得 map 35 左右，PB 能到 58.99 的代码流程 \n* 以上是训练 10 个 epoch 结果，需要大约2h；这里演示为了运行速度快一点， epoch 仅设了 3","metadata":{}},{"cell_type":"markdown","source":"# 安装工具包，下载字体、torch.vision 中 几个辅助工具类","metadata":{}},{"cell_type":"code","source":"# 安装 pycocotools\n!pip install pycocotools\n","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:36:55.627781Z","iopub.execute_input":"2021-08-01T03:36:55.628059Z","iopub.status.idle":"2021-08-01T03:37:02.137004Z","shell.execute_reply.started":"2021-08-01T03:36:55.628033Z","shell.execute_reply":"2021-08-01T03:37:02.136001Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 下载字体，用于显示目类型\n!rm -rf font\n!mkdir font\n!wget https://www.fontsquirrel.com/fonts/download/open-sans \n!unzip open-sans -d font","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:40:58.803881Z","iopub.execute_input":"2021-08-01T03:40:58.804225Z","iopub.status.idle":"2021-08-01T03:41:02.789095Z","shell.execute_reply.started":"2021-08-01T03:40:58.804192Z","shell.execute_reply":"2021-08-01T03:41:02.788058Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 下载 torch.vison 中几个工具类，主要训练和eval中方便记录参数信息\n!git clone https://github.com/pytorch/vision.git\n!cd vision\n!git checkout v0.3.0\n\n!cp vision/references/detection/utils.py ./\n!cp vision/references/detection/transforms.py ./\n!cp vision/references/detection/coco_eval.py ./\n!cp vision/references/detection/engine.py ./\n!cp vision/references/detection/coco_utils.py ./","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:41:07.024242Z","iopub.execute_input":"2021-08-01T03:41:07.024691Z","iopub.status.idle":"2021-08-01T03:41:12.535097Z","shell.execute_reply.started":"2021-08-01T03:41:07.024646Z","shell.execute_reply":"2021-08-01T03:41:12.533979Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!ls","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:41:16.10516Z","iopub.execute_input":"2021-08-01T03:41:16.105489Z","iopub.status.idle":"2021-08-01T03:41:16.745523Z","shell.execute_reply.started":"2021-08-01T03:41:16.105457Z","shell.execute_reply":"2021-08-01T03:41:16.744521Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!nvidia-smi","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:41:31.641448Z","iopub.execute_input":"2021-08-01T03:41:31.641831Z","iopub.status.idle":"2021-08-01T03:41:32.311399Z","shell.execute_reply.started":"2021-08-01T03:41:31.641795Z","shell.execute_reply":"2021-08-01T03:41:32.310204Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 参数定义","metadata":{}},{"cell_type":"code","source":"import torch\nimport torchvision\nfrom torch.utils.data import Dataset, DataLoader, Subset\n\nimport os\nimport numpy as np\nimport pandas as pd\n\nfrom PIL import Image, ImageDraw, ImageFont\n\nfrom pycocotools.coco import COCO\n\nimport matplotlib.pyplot as plt\n%matplotlib inline","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:41:34.714146Z","iopub.execute_input":"2021-08-01T03:41:34.714474Z","iopub.status.idle":"2021-08-01T03:41:36.006039Z","shell.execute_reply.started":"2021-08-01T03:41:34.714442Z","shell.execute_reply":"2021-08-01T03:41:36.005177Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CFG:\n    # dir\n    DATA_DIR = \"../input/cowboyoutfits/\"\n    TRAIN_JSON_PATH = os.path.join(DATA_DIR, \"train.json\")\n    TEST_CSV_PATH = os.path.join(DATA_DIR, \"test.csv\")\n    VALID_CSV_PATH = os.path.join(DATA_DIR, \"valid.csv\")\n    IMG_DIR = os.path.join(DATA_DIR, \"images\")\n    FONT_PATH = 'font/OpenSans-Semibold.ttf'\n\n    # dataset\n    BATCH_SIZE = 2\n    VALID_IMG_NUM = 500\n    NUM_WORKERS = 2\n    \n    \n    # train & model\n    NUM_CLASS = 6  # num_classes = class_num + 1\n    \n    # 为了运行快，这里用的 3，调到 10 提交后验证集 map 35 左右，榜单上大概能到 60\n    EPOCH = 15\n    PRINT_FREQ_TRAIN = 200\n    PRINT_FREQ_EVAL = 100\n    \n    # optim\n    LR = 0.005\n    WEIGHT_DECAY = 0.0005\n    MOMEMTUM = 0.9\n    \n    # lr_scheduler\n    STEP_SIZE = 5\n    GAMMA = 0.1\n    \n    # pre dict keys\n    KEY_SCORE = \"scores\"\n    KEY_LABELS = \"labels\"\n    KEY_BOXES = \"boxes\"","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:41:39.400244Z","iopub.execute_input":"2021-08-01T03:41:39.400583Z","iopub.status.idle":"2021-08-01T03:41:39.409043Z","shell.execute_reply.started":"2021-08-01T03:41:39.400533Z","shell.execute_reply":"2021-08-01T03:41:39.408111Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 用 coco 读取数据","metadata":{}},{"cell_type":"code","source":"coco = COCO(CFG.TRAIN_JSON_PATH)","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:41:53.771691Z","iopub.execute_input":"2021-08-01T03:41:53.772014Z","iopub.status.idle":"2021-08-01T03:41:53.951547Z","shell.execute_reply.started":"2021-08-01T03:41:53.771987Z","shell.execute_reply":"2021-08-01T03:41:53.950716Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 创建 id， 看下训练集中图片数量\nids = list(sorted(coco.imgs.keys()))\nlen(ids)","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:41:54.982407Z","iopub.execute_input":"2021-08-01T03:41:54.982794Z","iopub.status.idle":"2021-08-01T03:41:54.991747Z","shell.execute_reply.started":"2021-08-01T03:41:54.982761Z","shell.execute_reply":"2021-08-01T03:41:54.990917Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 创建用于类型映射的字典\ncat_map = {v['id']: v['name'] for k,v in coco.cats.items()}\n# 真实物体类别从1开始\ncatid_2_label = {cat_id: index + 1  for index, cat_id in enumerate(list(sorted(cat_map.keys())))}\nlabel_2_catid = {v: k for k, v in catid_2_label.items()}","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:41:56.538083Z","iopub.execute_input":"2021-08-01T03:41:56.539734Z","iopub.status.idle":"2021-08-01T03:41:56.549439Z","shell.execute_reply.started":"2021-08-01T03:41:56.539685Z","shell.execute_reply":"2021-08-01T03:41:56.548574Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"catid_2_label","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:41:57.622196Z","iopub.execute_input":"2021-08-01T03:41:57.622516Z","iopub.status.idle":"2021-08-01T03:41:57.627682Z","shell.execute_reply.started":"2021-08-01T03:41:57.622487Z","shell.execute_reply":"2021-08-01T03:41:57.626848Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 显示 2 张图片看看","metadata":{}},{"cell_type":"code","source":"# 显示 100 - 101 张图片\nsmp_ids = ids[100:100 + 2]\nsmp_file = coco.loadImgs(smp_ids)\nsmp_anno_ids = coco.getAnnIds(smp_ids)","metadata":{"execution":{"iopub.status.busy":"2021-07-31T16:12:15.370011Z","iopub.execute_input":"2021-07-31T16:12:15.370546Z","iopub.status.idle":"2021-07-31T16:12:15.375791Z","shell.execute_reply.started":"2021-07-31T16:12:15.370508Z","shell.execute_reply":"2021-07-31T16:12:15.374762Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 通过 coco，取出图片对应 anno 信息，画到图片上\nfont_size = 40\nfont = ImageFont.truetype(CFG.FONT_PATH, font_size)\nshow_imgs = []\n\nfor index in range(len(smp_ids)):\n    file_name = smp_file[index]['file_name']\n    file_path = os.path.join(CFG.IMG_DIR, file_name)\n    img = Image.open(file_path).convert(\"RGB\")\n    draw = ImageDraw.Draw(img)\n    \n    anno_id = smp_anno_ids[index]\n    targets = coco.loadAnns(anno_id)\n    for target in targets:\n        x_min, y_min, w, h = target['bbox']\n        x_max = x_min + w\n        y_max = y_min + h\n        \n        draw.rectangle((int(x_min), int(y_min), int(x_max), int(y_max)), outline='red', width=5)\n        \n        cat_name = cat_map[target['category_id']]\n        draw.text((int(x_min), int(y_min) - font_size - 10), cat_name, 'red', font, )\n        show_imgs.append(img)","metadata":{"execution":{"iopub.status.busy":"2021-07-31T16:12:15.377073Z","iopub.execute_input":"2021-07-31T16:12:15.377535Z","iopub.status.idle":"2021-07-31T16:12:15.552658Z","shell.execute_reply.started":"2021-07-31T16:12:15.377385Z","shell.execute_reply":"2021-07-31T16:12:15.551823Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.imshow(show_imgs[0])","metadata":{"execution":{"iopub.status.busy":"2021-07-31T16:12:15.553828Z","iopub.execute_input":"2021-07-31T16:12:15.554143Z","iopub.status.idle":"2021-07-31T16:12:15.841066Z","shell.execute_reply.started":"2021-07-31T16:12:15.55411Z","shell.execute_reply":"2021-07-31T16:12:15.840319Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 从 101 张图片可以看出，挺多太阳镜没有标出来\nplt.imshow(show_imgs[1])","metadata":{"execution":{"iopub.status.busy":"2021-07-31T16:12:15.843338Z","iopub.execute_input":"2021-07-31T16:12:15.843693Z","iopub.status.idle":"2021-07-31T16:12:16.099439Z","shell.execute_reply.started":"2021-07-31T16:12:15.843663Z","shell.execute_reply":"2021-07-31T16:12:16.098518Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 定义 Transforms、 DataSet 、DataLoader","metadata":{}},{"cell_type":"markdown","source":"## 引入 torch.vison github 中的 Transforms，能支持简单的基于目标检测的变换","metadata":{"execution":{"iopub.status.busy":"2021-07-31T08:17:54.59236Z","iopub.execute_input":"2021-07-31T08:17:54.592759Z","iopub.status.idle":"2021-07-31T08:17:54.597343Z","shell.execute_reply.started":"2021-07-31T08:17:54.592727Z","shell.execute_reply":"2021-07-31T08:17:54.596139Z"}}},{"cell_type":"code","source":"import transforms as T\ntrainforms_train = T.Compose([\n    T.ToTensor(),\n    \n    # 内部做了 bbox 变幻\n    T.RandomHorizontalFlip(0.5) \n    \n])\n\ntrainforms_test = T.Compose([\n    T.ToTensor()\n])","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:42:01.405927Z","iopub.execute_input":"2021-08-01T03:42:01.406254Z","iopub.status.idle":"2021-08-01T03:42:01.414291Z","shell.execute_reply.started":"2021-08-01T03:42:01.406225Z","shell.execute_reply":"2021-08-01T03:42:01.413376Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## CowBoyDataSet","metadata":{}},{"cell_type":"code","source":"class CowBoyDataSet(Dataset):\n    def __init__(self, coco, img_dir, transforms):\n        self.coco = coco\n        self.img_dir = img_dir\n        self.transforms = transforms\n        self.img_ids = list(sorted(coco.imgs.keys()))\n        \n    def __getitem__(self, idx):\n        # 返回 img tensor, bbox, cat_id\n        img_id = self.img_ids[idx]\n        img_name = self.coco.loadImgs(img_id)[0]['file_name']\n        img_path = os.path.join(self.img_dir, img_name)\n        img = Image.open(img_path).convert(\"RGB\")\n        \n        anno_ids = self.coco.getAnnIds(img_id)\n        annos = self.coco.loadAnns(anno_ids)\n        \n        boxes = []\n        labels = []\n        areas = []\n        iscrowds = []\n        \n        for anno in annos:\n            # anno 中 box 为 x,y,w,h\n            # vison 中 faster-rcnn box 为 x_min, y_min, x_max, y_max\n            x_min, y_min, w, h = anno['bbox']\n            x_max, y_max = x_min + w, y_min + h\n            boxes.append([x_min, y_min, x_max, y_max])\n            \n            cat_id = anno['category_id']\n            label = catid_2_label[cat_id]\n            labels.append(label)\n            areas.append(anno['area'])\n            \n            # 表示是否多个小物体聚在一起, false 为干净的单个物体\n            iscrowds.append(anno['iscrowd'])\n            \n        img_id = torch.as_tensor([idx], dtype=torch.int64)\n        boxes = torch.as_tensor(boxes, dtype=torch.float32)\n        labels = torch.as_tensor(labels, dtype=torch.int64)\n        areas = torch.tensor(areas)\n        iscrowds = torch.tensor(iscrowds)\n            \n        # 训练中 model 只用到 boxes 和 labels，其他参数是给 coco 工具 evalueate 时用的\n        targets = {\n            'boxes': boxes,\n            'labels': labels,\n            'image_id': img_id,\n            'area': areas,\n            'iscrowd': iscrowds\n        }\n            \n        if self.transforms is not None:\n            img, targets = self.transforms(img, targets)\n            \n        return img, targets\n            \n        \n        \n    def __len__(self):\n        return len(self.img_ids)","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:42:03.616526Z","iopub.execute_input":"2021-08-01T03:42:03.61693Z","iopub.status.idle":"2021-08-01T03:42:03.630242Z","shell.execute_reply.started":"2021-08-01T03:42:03.616895Z","shell.execute_reply":"2021-08-01T03:42:03.629245Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 划分训练、验证集，创建 dataloader","metadata":{}},{"cell_type":"code","source":"# 由于目标检测每张图片尺寸不一样，这个函数使得 dataloader 能过正确处理这种情况，否则会报错\ndef collate_fn(batch):\n    return tuple(zip(*batch))","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:42:08.938471Z","iopub.execute_input":"2021-08-01T03:42:08.938872Z","iopub.status.idle":"2021-08-01T03:42:08.944177Z","shell.execute_reply.started":"2021-08-01T03:42:08.938838Z","shell.execute_reply":"2021-08-01T03:42:08.942966Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 这里训练集、验证集，每次运行都是随机划分的，如果需要固定划分便于继续训练，需要保存 indices\ndataset_train = CowBoyDataSet(coco, CFG.IMG_DIR, trainforms_train)\ndataset_eval = CowBoyDataSet(coco, CFG.IMG_DIR, trainforms_test)\n\nindices = torch.randperm(len(dataset_train)).tolist()\ndataset_train = Subset(dataset_train, indices[: -CFG.VALID_IMG_NUM])\ndataset_eval = Subset(dataset_eval, indices[-CFG.VALID_IMG_NUM:])\n\ndataloader_train = DataLoader(\n    dataset_train, CFG.BATCH_SIZE, shuffle=True, num_workers=CFG.NUM_WORKERS,\n    collate_fn=collate_fn)\n\ndataloader_eval = DataLoader(\n    dataset_eval, CFG.BATCH_SIZE, shuffle=False, num_workers=CFG.NUM_WORKERS,\n    collate_fn=collate_fn)","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:42:30.780344Z","iopub.execute_input":"2021-08-01T03:42:30.780704Z","iopub.status.idle":"2021-08-01T03:42:30.810803Z","shell.execute_reply.started":"2021-08-01T03:42:30.780669Z","shell.execute_reply":"2021-08-01T03:42:30.810012Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 定义模型 - 基于 torch 自带的 faster-rcnn 和 retinanet 进下修改","metadata":{"execution":{"iopub.status.busy":"2021-07-31T08:22:29.754361Z","iopub.execute_input":"2021-07-31T08:22:29.754744Z","iopub.status.idle":"2021-07-31T08:22:29.758992Z","shell.execute_reply.started":"2021-07-31T08:22:29.754712Z","shell.execute_reply":"2021-07-31T08:22:29.758064Z"}}},{"cell_type":"markdown","source":"## torch 自带的 faster-rcnn，需要修改 box_predictor 以改变 num_classes","metadata":{"execution":{"iopub.status.busy":"2021-07-31T08:23:23.769553Z","iopub.execute_input":"2021-07-31T08:23:23.770018Z","iopub.status.idle":"2021-07-31T08:23:23.774408Z","shell.execute_reply.started":"2021-07-31T08:23:23.769952Z","shell.execute_reply":"2021-07-31T08:23:23.773421Z"}}},{"cell_type":"code","source":"def get_fast_rcnn_model():\n    model = fasterrcnn_resnet50_fpn(pretrained=True)\n    \n    in_features = model.roi_heads.box_predictor.cls_score.in_features\n    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, CFG.NUM_CLASS)\n    return model","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:44:34.979457Z","iopub.execute_input":"2021-08-01T03:44:34.979846Z","iopub.status.idle":"2021-08-01T03:44:34.984803Z","shell.execute_reply.started":"2021-08-01T03:44:34.979812Z","shell.execute_reply":"2021-08-01T03:44:34.983676Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 由于尝试 retinanet 效果较差，这个函数用于对比实验，验证这种weigth创建方式实分有问题，结果和第一种比没问题\ndef get_faster_rcnn_model_v2():\n    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\n    model_new = torchvision.models.detection.fasterrcnn_resnet50_fpn(\n        num_classes=CFG.NUM_CLASS)\n    \n    pretrained_state_dict = model.state_dict()\n    change_keys = [\n        'roi_heads.box_predictor.cls_score.weight',\n        'roi_heads.box_predictor.cls_score.bias',\n        'roi_heads.box_predictor.bbox_pred.weight',\n        'roi_heads.box_predictor.bbox_pred.bias'\n    ]\n    \n    \n    from_scratch_state_dict = model_new.state_dict()\n    \n   \n    for change_key in change_keys:\n        pretrained_state_dict[change_key] = from_scratch_state_dict[change_key]\n    \n    model_new.load_state_dict(pretrained_state_dict)\n    return model_new","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:44:35.423171Z","iopub.execute_input":"2021-08-01T03:44:35.423485Z","iopub.status.idle":"2021-08-01T03:44:35.429513Z","shell.execute_reply.started":"2021-08-01T03:44:35.423456Z","shell.execute_reply":"2021-08-01T03:44:35.428395Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# retinanet 使用更偷懒的方式修改 weight ，仅有受 num_classes 影响的层用默认初始化，其他层使用 pretrain 权重 \n# 存在问题，训练结果较差，用同样方式初始化 faster-rcnn 没问题，因此目前还未定位 retinanet 的问题在哪里\ndef get_retinanet_model():\n    model = torchvision.models.detection.retinanet_resnet50_fpn(pretrained=True)\n    model_new = torchvision.models.detection.retinanet_resnet50_fpn(\n        num_classes=CFG.NUM_CLASS)\n    \n    pretrained_state_dict = model.state_dict()\n    change_keys = [\n        'head.classification_head.cls_logits.weight',\n        'head.classification_head.cls_logits.bias'\n    ]\n    \n    \n    from_scratch_state_dict = model_new.state_dict()\n    \n   \n    for change_key in change_keys:\n        pretrained_state_dict[change_key] = from_scratch_state_dict[change_key]\n    \n    model_new.load_state_dict(pretrained_state_dict)\n    return model_new","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:44:36.620861Z","iopub.execute_input":"2021-08-01T03:44:36.6212Z","iopub.status.idle":"2021-08-01T03:44:36.627863Z","shell.execute_reply.started":"2021-08-01T03:44:36.621171Z","shell.execute_reply":"2021-08-01T03:44:36.626757Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 创建模型、优化器、lr_scheduler","metadata":{}},{"cell_type":"code","source":"model = get_retinanet_model()\ndevice = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\nparam = [p for p in model.parameters() if p.requires_grad]\noptmizer = torch.optim.SGD(param, lr=CFG.LR, momentum=CFG.MOMEMTUM)\nlr_scheduler = torch.optim.lr_scheduler.StepLR(optmizer, step_size=CFG.STEP_SIZE, gamma=CFG.GAMMA)","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:44:39.165783Z","iopub.execute_input":"2021-08-01T03:44:39.166108Z","iopub.status.idle":"2021-08-01T03:44:56.109542Z","shell.execute_reply.started":"2021-08-01T03:44:39.166079Z","shell.execute_reply":"2021-08-01T03:44:56.108622Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 训练","metadata":{}},{"cell_type":"markdown","source":"## train model","metadata":{"execution":{"iopub.status.busy":"2021-07-31T08:31:13.433682Z","iopub.execute_input":"2021-07-31T08:31:13.434077Z","iopub.status.idle":"2021-07-31T08:31:13.438723Z","shell.execute_reply.started":"2021-07-31T08:31:13.434042Z","shell.execute_reply":"2021-07-31T08:31:13.437451Z"}}},{"cell_type":"code","source":"# 由于faster-rcnn 已经封装好了 cls 和 bbox 的 loss，其实自己写一个 train 函数很简单，和图片分类几乎一样\n# 这里使用从 github: torch.vison 拷贝的工具类，主要是能详细打印训练/验证时多个关键信息，并且组织较美观\n# 还有一点影响训练的是，工具类中使用了另一个 warmup scheduler，以降低训练初始学习率\nfrom engine import train_one_epoch, evaluate\nimport utils","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:47:04.162621Z","iopub.execute_input":"2021-08-01T03:47:04.163008Z","iopub.status.idle":"2021-08-01T03:47:04.178767Z","shell.execute_reply.started":"2021-08-01T03:47:04.162975Z","shell.execute_reply":"2021-08-01T03:47:04.177784Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.to(device)","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:47:05.726932Z","iopub.execute_input":"2021-08-01T03:47:05.72725Z","iopub.status.idle":"2021-08-01T03:47:09.786412Z","shell.execute_reply.started":"2021-08-01T03:47:05.727221Z","shell.execute_reply":"2021-08-01T03:47:09.78563Z"},"collapsed":true,"jupyter":{"outputs_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!mkdir train_state","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 这里为了速度快，演示代码只训练 3 epoch（半小时左右执行完），达到 10 epoch（约2小时） 85+ 水平\n# 这会导致我使用的 lr_scheduler 未生效（每 3epoch 降低 90% 学习率），然后保存模型训练状态，方便后续继续训练\n# 如果用于比赛提交，可以训练 10+ epoch\nfor epoch in range(CFG.EPOCH):\n    # train\n    train_one_epoch(model, optmizer, dataloader_train, device, epoch, CFG.PRINT_FREQ_TRAIN)\n    lr_scheduler.step()\n    \n    # eval\n    evaluate(model, dataloader_eval, device)\n    \n    # save\n    if epoch %3 == 0 and epoch > 1:\n      train_state_dict = {\n          'model':model.state_dict(),\n          'optim':optmizer.state_dict(),\n          'lr_scheduler':lr_scheduler.state_dict(),\n          'indices':indices,\n          'info':'cow_bow_retinanet-epoch' + str(epoch)\n      }\n      torch.save(train_state_dict, \n                 'train_state/retinanet-epoch' + str(epoch))","metadata":{"execution":{"iopub.status.busy":"2021-08-01T03:51:50.214992Z","iopub.execute_input":"2021-08-01T03:51:50.215381Z","iopub.status.idle":"2021-08-01T04:33:58.440067Z","shell.execute_reply.started":"2021-08-01T03:51:50.215349Z","shell.execute_reply":"2021-08-01T04:33:58.438902Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 评估指标","metadata":{}},{"cell_type":"markdown","source":"* evaluate 的第一行，即为比赛所用的评估指标：\n* Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.261","metadata":{}},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 保存训练状态","metadata":{}},{"cell_type":"code","source":"# 保存训练状态，以便随时中断、后续继续训练\ntrain_state_dict = {\n    'model':model.state_dict(),\n    'optim':optmizer.state_dict(),\n    'lr_scheduler':lr_scheduler.state_dict(),\n    'indices':indices,\n    'info':'cow_bow_fascter-rcnn-epoch2'\n}\n\ntorch.save(\n    train_state_dict, \n    'train_state_retinanet_epoch' + str(CFG.EPOCH))","metadata":{"execution":{"iopub.status.busy":"2021-08-01T04:44:41.232025Z","iopub.execute_input":"2021-08-01T04:44:41.232377Z","iopub.status.idle":"2021-08-01T04:44:42.183367Z","shell.execute_reply.started":"2021-08-01T04:44:41.232345Z","shell.execute_reply":"2021-08-01T04:44:42.182527Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 预测","metadata":{}},{"cell_type":"markdown","source":"## 预测用的 dataset 和 dataloader","metadata":{}},{"cell_type":"code","source":"class CowBoyPredictDataSet(Dataset):\n    def __init__(self, csv_path, img_dir, transforms):\n        self.csv_path = csv_path\n        self.img_dir = img_dir\n        self.transforms = transforms\n        self.df_predict = pd.read_csv(self.csv_path)\n        \n    def __getitem__(self, idx):\n        # 返回 img tensor\n        row = self.df_predict.iloc[idx]\n        row_id = row['id']\n        img_name = row['file_name']\n        \n        img_path = os.path.join(self.img_dir, img_name)\n        img = Image.open(img_path).convert(\"RGB\")\n        \n            \n        if self.transforms is not None:\n            img, _ = self.transforms(img, None)\n            \n        return img, row_id\n            \n        \n        \n    def __len__(self):\n        return self.df_predict.shape[0]","metadata":{"execution":{"iopub.status.busy":"2021-08-01T04:44:45.401286Z","iopub.execute_input":"2021-08-01T04:44:45.401619Z","iopub.status.idle":"2021-08-01T04:44:45.408799Z","shell.execute_reply.started":"2021-08-01T04:44:45.401588Z","shell.execute_reply":"2021-08-01T04:44:45.407873Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dataset_predict = CowBoyPredictDataSet(CFG.VALID_CSV_PATH, CFG.IMG_DIR, trainforms_test)\n\ndataloader_predict = DataLoader(\n    dataset_predict, 1, shuffle=False, num_workers=CFG.NUM_WORKERS,\n    collate_fn=collate_fn)","metadata":{"execution":{"iopub.status.busy":"2021-08-01T04:44:46.369355Z","iopub.execute_input":"2021-08-01T04:44:46.369687Z","iopub.status.idle":"2021-08-01T04:44:46.391411Z","shell.execute_reply.started":"2021-08-01T04:44:46.369656Z","shell.execute_reply":"2021-08-01T04:44:46.390638Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 看几张预测结果","metadata":{}},{"cell_type":"code","source":"font_size = 40\nfont = ImageFont.truetype(CFG.FONT_PATH, font_size)\ntramsform_pil = torchvision.transforms.ToPILImage()\n\n\ndef see_pre_result(index, device='cpu', threshold = 0.5):\n    img_tensor, image_ids = dataset_predict[index]\n    img = tramsform_pil(img_tensor)\n    \n    model.eval()\n    with torch.no_grad():\n        img_tensor = img_tensor.unsqueeze(0)\n        img_tensor = img_tensor.to(device)\n        pre_batch = model(img_tensor)\n\n    \n        \n    pre = pre_batch[0]\n    print(pre['labels'])\n    print(pre['scores'])\n    \n    draw = ImageDraw.Draw(img)\n    length = pre['labels'].shape[0]\n    for i in range(length):\n        score = pre['scores'][i].item()\n        if score < threshold:\n            continue\n        \n        x_min, y_min, x_max, y_max = pre['boxes'][i]\n        x_min, y_min, x_max, y_max = int(x_min), int(y_min), int(x_max), int(y_max)\n\n        label_id = pre['labels'][i].item()\n        cat_name = cat_map[label_2_catid[label_id]]\n        color = (100 + 30 * label_id, 0, 0)\n\n        draw.rectangle((x_min, y_min, x_max, y_max), outline=color, width=5)\n        draw.text((x_min, y_min - font_size), cat_name, color, font)\n\n    plt.imshow(img)\n    \n    return pre_batch, image_ids","metadata":{"execution":{"iopub.status.busy":"2021-08-01T04:44:48.509778Z","iopub.execute_input":"2021-08-01T04:44:48.510119Z","iopub.status.idle":"2021-08-01T04:44:48.521618Z","shell.execute_reply.started":"2021-08-01T04:44:48.510091Z","shell.execute_reply":"2021-08-01T04:44:48.520513Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pre_batch = see_pre_result(100, device)","metadata":{"execution":{"iopub.status.busy":"2021-08-01T04:44:49.378889Z","iopub.execute_input":"2021-08-01T04:44:49.379216Z","iopub.status.idle":"2021-08-01T04:44:49.829101Z","shell.execute_reply.started":"2021-08-01T04:44:49.379185Z","shell.execute_reply":"2021-08-01T04:44:49.828135Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pre_batch = see_pre_result(500, device)","metadata":{"execution":{"iopub.status.busy":"2021-08-01T04:45:26.509004Z","iopub.execute_input":"2021-08-01T04:45:26.509321Z","iopub.status.idle":"2021-08-01T04:45:27.110419Z","shell.execute_reply.started":"2021-08-01T04:45:26.509293Z","shell.execute_reply":"2021-08-01T04:45:27.10508Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pre_batch = see_pre_result(700, device)","metadata":{"execution":{"iopub.status.busy":"2021-08-01T04:45:34.607051Z","iopub.execute_input":"2021-08-01T04:45:34.607401Z","iopub.status.idle":"2021-08-01T04:45:35.021246Z","shell.execute_reply.started":"2021-08-01T04:45:34.60737Z","shell.execute_reply":"2021-08-01T04:45:35.020251Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 结果转换为 json","metadata":{}},{"cell_type":"code","source":"# 进行预测\npre_results = []\nmodel.eval()\nwith torch.no_grad():\n  for img_tensor, image_ids in dataloader_predict:\n    img_tensor = img_tensor[0]\n    image_ids = image_ids[0]\n    img_tensor = img_tensor.unsqueeze(0)\n    img_tensor = img_tensor.to(device)\n    pre_batch = model(img_tensor)\n    pre_batch = pre_batch\n    pre_results.append((pre_batch, image_ids))","metadata":{"execution":{"iopub.status.busy":"2021-08-01T04:45:39.651313Z","iopub.execute_input":"2021-08-01T04:45:39.651658Z","iopub.status.idle":"2021-08-01T04:47:14.605686Z","shell.execute_reply.started":"2021-08-01T04:45:39.651628Z","shell.execute_reply":"2021-08-01T04:47:14.604453Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 模型输出，转化为提交用的 dict 的函数，方便后续直接 josn dump\n# ！注意：这个函数写的简单，仅支持预测时 batch size 为 1\ndef out_to_dict(pre_results, threshold=0.5):\n    results = []\n    for one_pre in pre_results:\n        image_id = int(one_pre[1])\n        \n        # 需要固定预测试是，batch size 为 1，否则出错\n        data_dict = one_pre[0][0]\n        if CFG.KEY_SCORE not in data_dict:\n            continue\n            \n        pre_num = len(data_dict[CFG.KEY_SCORE])\n        for i in range(pre_num):\n            score = data_dict[CFG.KEY_SCORE][i].to('cpu').item()\n            if score < threshold:\n                continue\n                \n            label_id = int(data_dict[CFG.KEY_LABELS][i].to('cpu').item())\n            bbox = data_dict[CFG.KEY_BOXES][i].to('cpu').tolist()\n            w = bbox[2] - bbox[0]\n            h = bbox[3] - bbox[1]\n            bbox[2] = w\n            bbox[3] = h\n            out = {\n                'image_id': image_id,\n                'category_id': label_2_catid[label_id],\n                'bbox': bbox,\n                'score': score\n            }\n            results.append(out)\n            \n    return results","metadata":{"execution":{"iopub.status.busy":"2021-08-01T04:47:14.607355Z","iopub.execute_input":"2021-08-01T04:47:14.607764Z","iopub.status.idle":"2021-08-01T04:47:14.616518Z","shell.execute_reply.started":"2021-08-01T04:47:14.607729Z","shell.execute_reply":"2021-08-01T04:47:14.615672Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 拿一张看看，主要检查提交用文件的坐标问题","metadata":{}},{"cell_type":"code","source":"# 检查 bbox， 模型输出是 x_min, y_min, x_max, y_max, 上传需要 x_min, y_min, w, h\nresult_dict = out_to_dict(pre_results)","metadata":{"execution":{"iopub.status.busy":"2021-08-01T04:47:14.618913Z","iopub.execute_input":"2021-08-01T04:47:14.619303Z","iopub.status.idle":"2021-08-01T04:47:15.042375Z","shell.execute_reply.started":"2021-08-01T04:47:14.619266Z","shell.execute_reply":"2021-08-01T04:47:15.041316Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 看一行用于提交 bbox 的结果，这并不是一张图片上所有预测出的 bbox，因为按提交格式一张图片多个 bbox 会分多行输出\ndf_valid = pd.read_csv(CFG.VALID_CSV_PATH)\ndata1 = result_dict[1]\nimg1_id = data1['image_id']\nfile1 = os.path.join(CFG.IMG_DIR, df_valid[df_valid.id == img1_id].iloc[0].file_name)\nimg1 = Image.open(file1).convert(\"RGB\")\n\ndraw = ImageDraw.Draw(img1)\nx_min, y_min, w, h = data1['bbox']\nx_max = x_min + w\ny_max = y_min + h\n\ndraw.rectangle((int(x_min), int(y_min), int(x_max), int(y_max)), outline='red', width=5)\n\ncat_name = cat_map[data1['category_id']]\ndraw.text((int(x_min), int(y_min) - font_size), cat_name, 'red', font, )\nplt.imshow(img1)","metadata":{"execution":{"iopub.status.busy":"2021-08-01T04:47:15.043888Z","iopub.execute_input":"2021-08-01T04:47:15.044246Z","iopub.status.idle":"2021-08-01T04:47:15.338034Z","shell.execute_reply.started":"2021-08-01T04:47:15.044207Z","shell.execute_reply":"2021-08-01T04:47:15.337202Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 将 dict dump 为提交所用文件 ","metadata":{}},{"cell_type":"code","source":"import json\n# 输出要求 zip 文件， zip 中保护一个叫 answer.json 的文件\nout_path = 'answer.json'\nwith open(out_path, 'w') as f:\n  json.dump(result_dict, f)","metadata":{"execution":{"iopub.status.busy":"2021-08-01T04:47:15.339237Z","iopub.execute_input":"2021-08-01T04:47:15.339724Z","iopub.status.idle":"2021-08-01T04:47:15.40233Z","shell.execute_reply.started":"2021-08-01T04:47:15.339684Z","shell.execute_reply":"2021-08-01T04:47:15.401385Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}