Run_PointTree_S3DIS.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "[Release]Run_PointTree_S3DIS.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "xJm_DG6kzGYG",
        "4lmXaxD5ynlo",
        "oLvduaqgSRuq",
        "-LzQmMFZmz60",
        "rqM9d9ravexM",
        "0CsaatrcpXKk"
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4lmXaxD5ynlo"
      },
      "source": [
        "# **Starting Work**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X7mHJSTOLfAk"
      },
      "source": [
        "from os import chdir, environ"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!git pull"
      ],
      "metadata": {
        "id": "NfavCQzG8-qP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IHDHoF-Lsmle"
      },
      "source": [
        "!ls -lt --time-style='+%y-%m-%d %H:%M:%S'\n",
        "!dir"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oLvduaqgSRuq"
      },
      "source": [
        "# **Experiment Init**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XInyKWaTSYXY"
      },
      "source": [
        "import json\n",
        "import h5py\n",
        "import torch\n",
        "import os\n",
        "from encoder import Encoder, MLP\n",
        "from segment import Segment\n",
        "from res_encoder import ResSegment\n",
        "from dataset import *\n",
        "# from encoders.model_cross import Encoder, init_logging, MLP\n",
        "from build_tree import get_directions, init_directions\n",
        "import logging\n",
        "\n",
        "OUTPUT = 'scratch_s3dis'\n",
        "\n",
        "#encoder\n",
        "model_size = 4096\n",
        "sample_layers = 50 # 2\n",
        "channel = 1\n",
        "dim = 4096 # 2048\n",
        "dim_layer0 = 16 # 16\n",
        "dim_repeat_cut = 5\n",
        "use_sym = False #True\n",
        "use_tnet = True # True\n",
        "\n",
        "#decoder \n",
        "ancestor_dim = 512 # 512\n",
        "model_name = \"segment\"\n",
        "use_dyn_tree = False\n",
        "\n",
        "\n",
        "#homo w/o pa\n",
        "# prefix = \"_homo_nopa_xinf\"\n",
        "# transform = homo_transform\n",
        "# augment = 1\n",
        "# no_prealign = True \n",
        "# rotate_only = False \n",
        "# load_pretrain = False\n",
        "# augment_fn = lambda pts: augment_generator(pts) #, fetch_perm=True), dropout=0.5) #, shift=True) #, scale=True, agg_coef=0.5)\n",
        "# pca_augment = False\n",
        "# test_as_valid = True #True\n",
        "# trunc = 6\n",
        "# part_cls_dropout = None\n",
        "\n",
        "\n",
        "# #homo\n",
        "# prefix = \"_homo_xinf\"\n",
        "# transform = homo_transform\n",
        "# augment = 1\n",
        "# no_prealign = False \n",
        "# rotate_only = False \n",
        "# load_pretrain = False\n",
        "# augment_fn = lambda pts: augment_generator(pts) #, fetch_perm=True), dropout=0.5) #, shift=True) #, scale=True, agg_coef=0.5)\n",
        "# pca_augment = False\n",
        "# test_as_valid = True #True\n",
        "# trunc = 6\n",
        "# part_cls_dropout = None\n",
        "\n",
        "# # orig\n",
        "# prefix = \"_orig_xinf\"\n",
        "# transform = no_transform\n",
        "# augment = 1\n",
        "# no_prealign = False #True \n",
        "# rotate_only = True #False\n",
        "# load_pretrain = False\n",
        "# augment_fn = lambda pts: augment_generator(pts, fetch_perm=True, dropout=0.5) #, shift=True) #, scale=True, agg_coef=0.5)\n",
        "# pca_augment = False\n",
        "# test_as_valid = True #True\n",
        "# trunc = 9\n",
        "# part_cls_dropout = 0.5 #0.5\n",
        "\n",
        "# #affine\n",
        "prefix = \"_affine_xinf\"\n",
        "transform = affine_transform\n",
        "augment = 1\n",
        "no_prealign = False \n",
        "rotate_only = False \n",
        "load_pretrain = False\n",
        "augment_fn = lambda pts: augment_generator(pts) #, fetch_perm=True), dropout=0.5) #, shift=True) #, scale=True, agg_coef=0.5)\n",
        "pca_augment = False\n",
        "test_as_valid = True #True\n",
        "trunc = 6\n",
        "part_cls_dropout = None\n",
        "\n",
        "# #affine w/o pa\n",
        "# prefix = \"_affine_nopa_xinf\"\n",
        "# transform = affine_transform\n",
        "# augment = 1\n",
        "# no_prealign = True \n",
        "# rotate_only = False \n",
        "# load_pretrain = False\n",
        "# augment_fn = lambda pts: augment_generator(pts) #, fetch_perm=True), dropout=0.5) #, shift=True) #, scale=True, agg_coef=0.5)\n",
        "# pca_augment = False\n",
        "# test_as_valid = True #True\n",
        "# trunc = 6\n",
        "# part_cls_dropout = None\n",
        "\n",
        "\n",
        "sample_child_first = False # True in l7s1\n",
        "num_parts = 13\n",
        "test_area = 5\n",
        "DATASET = './datasets/S3DIS_hdf5'\n",
        "chaos_limit = 0\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S9Rt8bf5ScUW"
      },
      "source": [
        "global logging_init_flag\n",
        "logging_init_flag = False\n",
        "\n",
        "def init_logging(OUTPUT):\n",
        "    global logging_init_flag\n",
        "    if logging_init_flag is True:\n",
        "        return\n",
        "    logging_init_flag = True\n",
        "\n",
        "    logger = logging.getLogger()\n",
        "    logger.setLevel(logging.DEBUG)\n",
        "    formatter = logging.Formatter(\n",
        "        '%(asctime)s - %(levelname)s:\\t%(message)s',\n",
        "        datefmt='%Y-%m-%d %H:%M:%S')\n",
        "\n",
        "    fh = logging.FileHandler(f\"{OUTPUT}/training.log\")\n",
        "    fh.setLevel(logging.INFO)\n",
        "    fh.setFormatter(formatter)\n",
        "\n",
        "    ch = logging.StreamHandler()\n",
        "    ch.setLevel(logging.DEBUG)\n",
        "    ch.setFormatter(formatter)\n",
        "\n",
        "    logger.addHandler(ch)\n",
        "    logger.addHandler(fh)\n",
        "\n",
        "init_logging(OUTPUT)\n",
        "logging.info(f\"prefix = {prefix}\")\n",
        "_ = init_directions(chaos_limit, calc_dmap=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-ve1aug-tytB"
      },
      "source": [
        "torch.manual_seed(674433238)\n",
        "def new_model(num_parts=num_parts):\n",
        "    if model_name == 'res_segment':\n",
        "        def encoder_fn(point_dim):\n",
        "            encoder = Encoder(model_size, sample_layers, dim, OUTPUT, point_dim=point_dim, extra_dim=trunc - 3, channel=channel, sample_child_first=sample_child_first, dim_layer0=dim_layer0, layer0_mlp_dim=ancestor_dim, dim_repeat_cut=dim_repeat_cut, use_symmetry_loss=use_sym).cuda()\n",
        "            if not use_tnet:\n",
        "                encoder.layers[0].pts_align = torch.nn.Identity()\n",
        "            return encoder\n",
        "        model = ResSegment(ancestor_dim, encoder_fn, use_dyn_tree=use_dyn_tree, carry_dim_seg1=min(dim//16, ancestor_dim), num_parts=num_parts, part_cls_dropout=part_cls_dropout).cuda()   \n",
        "    else:\n",
        "        encoder = Encoder(model_size, sample_layers, dim, OUTPUT, channel=channel, extra_dim=trunc - 3, sample_child_first=sample_child_first, dim_layer0=dim_layer0, dim_repeat_cut=dim_repeat_cut, use_symmetry_loss=use_sym).cuda()\n",
        "        if not use_tnet:\n",
        "            encoder.layers[0].pts_align = torch.nn.Identity()\n",
        "        model = Segment(ancestor_dim, encoder, num_parts=num_parts, part_cls_dropout=part_cls_dropout).cuda()\n",
        "    model.tree.use_sym = False\n",
        "    return model\n",
        "\n",
        "model = new_model()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qS58zLcFNjT9"
      },
      "source": [
        "model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rqM9d9ravexM"
      },
      "source": [
        "# **Data**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KWRjuhs4tl1v"
      },
      "source": [
        "from dataset import *\n",
        "import numpy as np\n",
        "from math import ceil\n",
        "from tqdm import tqdm"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AqqwreNL5L6g"
      },
      "source": [
        "make = make_data_default\n",
        "if rotate_only:\n",
        "    make = make_data_rotate_only\n",
        "if no_prealign:\n",
        "    make = make_data_no_prealign"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "clouds = []\n",
        "labels = []\n",
        "\n",
        "for i in tqdm(range(24)):\n",
        "    data_file = h5py.File(f'{DATASET}/ply_data_all_{i}.h5')\n",
        "    clouds.append(torch.tensor(np.array(data_file['data'])))\n",
        "    labels.append(torch.tensor(np.array(data_file['label'])))\n",
        "\n",
        "clouds = torch.cat(clouds, dim=0)\n",
        "labels = torch.cat(labels, dim=0)"
      ],
      "metadata": {
        "id": "4x2ZvLNYk2UU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_ind = []\n",
        "test_ind = []\n",
        "\n",
        "with open(f\"{DATASET}/room_filelist.txt\") as file:\n",
        "    for i, filename in enumerate(file.read().split()):\n",
        "        (test_ind if filename.startswith(f\"Area_{test_area}_\") else train_ind).append(i)\n",
        "\n",
        "train_ind = torch.tensor(train_ind)\n",
        "test_ind = torch.tensor(test_ind)"
      ],
      "metadata": {
        "id": "y4-hlEpalazY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vG7rbo4ekcjE"
      },
      "source": [
        "train_dataset = PointCloudDataset(clouds, labels, model.tree.arrange, trunc=trunc, use_norm=True, subset=train_ind, augment=augment, transform=transform, make=make, augment_fn=augment_fn, sample_points=model_size)\n",
        "test_dataset = PointCloudDataset(clouds, labels, model.tree.arrange, trunc=trunc, use_norm=True, subset=test_ind, augment=1, transform=transform, make=make, sample_points=model_size)\n",
        "valid_dataset = test_dataset"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bX9UCfSsTweO"
      },
      "source": [
        "len(train_dataset), len(valid_dataset), len(test_dataset)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oiILxo2qlfBf"
      },
      "source": [
        "# **Train**"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Defs"
      ],
      "metadata": {
        "id": "0CsaatrcpXKk"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nw0-l-jDmbMO"
      },
      "source": [
        "from dataset import make_batch_train, make_batch_eval\n",
        "import torch.nn as nn\n",
        "\n",
        "global batch_size\n",
        "batch_size = 64\n",
        "\n",
        "for name, dataset in zip(['train', 'valid', 'test'], [train_dataset, valid_dataset, test_dataset]):\n",
        "    print(f\"Loading {name}\")\n",
        "    dataset.mem = None\n",
        "\n",
        "    try:\n",
        "        tmp = torch.load(f'{OUTPUT}/{name}_data{prefix}.pth')\n",
        "        dataset.mem = tmp\n",
        "        assert len(dataset.mem) == len(dataset), f'Size unmatch: {len(dataset.mem)} != {len(dataset)}'\n",
        "        continue\n",
        "    except:\n",
        "        print(\"Try part mode\")\n",
        "\n",
        "    dataset.mem = []\n",
        "    for i in range(0, 1000000000):\n",
        "        try:\n",
        "            tmp = torch.load(f'{OUTPUT}/{name}_data{prefix}.{i}.pth')\n",
        "            dataset.mem += tmp\n",
        "            print(f\"Loaded part {i} # = {len(tmp)}\")\n",
        "        except:\n",
        "            break\n",
        "    if len(dataset.mem) > 0:\n",
        "        continue\n",
        "\n",
        "    print(\"Use force online\")\n",
        "    dataset.mem = None\n",
        "    dataset.force_online = True\n",
        "\n",
        "def inf_iter(a):\n",
        "    while True:\n",
        "        for k in a:\n",
        "            yield(k)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5Dj6kfNZmdeq"
      },
      "source": [
        "class AccStat:\n",
        "    def __init__(self):\n",
        "        self.clear()\n",
        "\n",
        "    def clear(self):\n",
        "        self.correct = 0\n",
        "        self.total = 0\n",
        "\n",
        "    def add(self, cor, num=1, mean=False):\n",
        "        if mean:\n",
        "            cor *= num\n",
        "        self.correct += cor\n",
        "        self.total += num\n",
        "\n",
        "    def result(self, clear=False):\n",
        "        ret = self.correct / max(self.total, 1e-5)\n",
        "        if clear:\n",
        "            self.clear()\n",
        "        return ret \n",
        "\n",
        "    def __str__(self):\n",
        "        return \"%.4lf\" % self.result()\n",
        "\n",
        "def calc_miou(logits, part_label):\n",
        "    if len(logits.shape) < 3:\n",
        "        logits = logits.unsqueeze(0)\n",
        "    \n",
        "    if len(part_label.shape) < 2:\n",
        "        part_label = part_label.unsqueeze(0)\n",
        "\n",
        "    pred = logits.cuda().argmax(dim=-1)\n",
        "    part = torch.arange(logits.size(-1)).cuda()\n",
        "\n",
        "    in_pred = part[None, :, None] == pred[:, None, :]\n",
        "    in_label = part[None, :, None] == part_label.cuda()[:, None, :]\n",
        "\n",
        "    I = (in_pred & in_label).sum(dim=-1)\n",
        "    U = (in_pred | in_label).sum(dim=-1)\n",
        "\n",
        "    part_IOU = I / U.clamp(min=1).float()\n",
        "    part_IOU[U < 0.5] = 1\n",
        "    shape_IOU = part_IOU.mean(dim=-1)\n",
        "    return shape_IOU.mean().item()\n",
        "\n",
        "def evaluate(model, linear, loader, noprint=False, perms=[None]):\n",
        "    cls = AccStat()\n",
        "    stat = AccStat()\n",
        "    miou = AccStat()\n",
        "\n",
        "    if not noprint:\n",
        "        logging.info(f\"loader # = {len(loader)}\")\n",
        "\n",
        "    print_epoch = 1\n",
        "\n",
        "    model.eval()\n",
        "    linear.eval()\n",
        "    activate = lambda x : x\n",
        "\n",
        "    for epoch, (input, part_label) in enumerate(loader):\n",
        "        part_label = part_label.cuda()\n",
        "        with torch.no_grad():\n",
        "            for iperm in perms:\n",
        "                features = linear(model(*input, perm=iperm))\n",
        "                logits_all = model.part_classfier(features)\n",
        "\n",
        "                stat.add((logits_all.max(dim=-1).indices == part_label).float().mean())\n",
        "                miou.add(calc_miou(logits_all, part_label))\n",
        "\n",
        "                   \n",
        "        if not noprint:\n",
        "            if (epoch // batch_size + 1) % print_epoch == 0:\n",
        "                logging.debug(f\"test #{epoch} correct = {'%.6lf' % stat.result()}\")\n",
        "\n",
        "    if not noprint:\n",
        "        logging.info(f\"Done: score = {'%.8lf' % stat.result()}\")\n",
        "\n",
        "    model.train()\n",
        "    linear.train()  \n",
        "    \n",
        "    return stat.result(), \"%.6lf %.6lf\" % (stat.result(), miou.result())\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch.nn.functional as F\n",
        "def train(more_epoch=100000, current_epoch=0, ckpt=None, valid_result_threshold=999.0):\n",
        "    best_vres = -1.0\n",
        "    \n",
        "    from random import choice, randint\n",
        "    import build_tree\n",
        "    import torch\n",
        "    \n",
        "    model = new_model(num_parts).cuda()\n",
        "    linear = nn.Identity()\n",
        "    model.train()\n",
        "    linear.train()\n",
        "    activate = lambda x : x\n",
        "\n",
        "    def get_trans(n=3):\n",
        "        if n == 0:  return [None]\n",
        "        return [randint(0, len(build_tree.transforms) - 1) for _ in range(n)]\n",
        "\n",
        "    global batch_size \n",
        "\n",
        "    num_workers = 12\n",
        "    mbtrain = make_batch_generator(pca_augment=pca_augment and not no_prealign)\n",
        "    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, collate_fn=mbtrain, pin_memory=True, drop_last=True)\n",
        "    valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, collate_fn=make_batch_eval, pin_memory=True, drop_last=False)\n",
        "    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, collate_fn=make_batch_eval, pin_memory=True, drop_last=False)\n",
        "\n",
        "    logging.info(f\"train epoch = {current_epoch + 1} ~ {more_epoch} threshold = {valid_result_threshold} count = {len(train_loader) * batch_size} num_parts = {num_parts}\")\n",
        "\n",
        "    cum_loss = 0\n",
        "    cum_inner_loss = 0\n",
        "\n",
        "    batch_scale = 1 # 1 if basic else 1\n",
        "    epoch_scale = max(1, 256 // batch_size) # 8 # 4 if basic else 2\n",
        "    num_trans = 0 #0 if no_prealign else 1\n",
        "\n",
        "\n",
        "    print_epoch = 20\n",
        "    valid_epoch = 100\n",
        "    epoch_since = 0\n",
        "    save_epoch = 100\n",
        "    cut_epoch = 10000000000\n",
        "\n",
        "    cloud_accuracy = AccStat()\n",
        "    accuracy = AccStat()\n",
        "    miou = AccStat()\n",
        "    \n",
        "    cloud_classify_coef = 0.3\n",
        "    part_classify_coef = 0.2\n",
        "    class_spec_coef = 1\n",
        "\n",
        "    threshold = -1.0\n",
        "\n",
        "    temperature = 1 # 0.07\n",
        "\n",
        "    crit = torch.nn.CrossEntropyLoss(ignore_index=-1)\n",
        "    opt = torch.optim.Adam(list(model.parameters()) + list(linear.parameters()), lr=1e-4)\n",
        "    # sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=num_epoch / 5, eta_min=1e-5)\n",
        "    # sch = torch.optim.lr_scheduler.ExponentialLR(opt, 0.998)\n",
        "    sch = torch.optim.lr_scheduler.ExponentialLR(opt, 0.9999)\n",
        "\n",
        "    if ckpt is not None:\n",
        "        filename = ckpt\n",
        "        ckpt = torch.load(ckpt)\n",
        "        model.load_state_dict(ckpt['model'])\n",
        "        linear.load_state_dict(ckpt['linear'])\n",
        "        opt.load_state_dict(ckpt['opt'])\n",
        "        sch.load_state_dict(ckpt['sch'])\n",
        "        best_vres = ckpt['best_vres']\n",
        "        logging.info(f\"Loaded ckpt {filename} best_vres = {best_vres}\")\n",
        "\n",
        "    def save(epoch):\n",
        "        torch.save({\n",
        "            'model': model.state_dict(),\n",
        "            'linear': linear.state_dict(),\n",
        "            'opt': opt.state_dict(),\n",
        "            'sch': sch.state_dict(),\n",
        "            'best_vres': best_vres,\n",
        "        }, f\"{OUTPUT}/trained_{epoch}.pth\")\n",
        "\n",
        "\n",
        "    save(0)\n",
        "\n",
        "    train_iter = inf_iter(train_loader)\n",
        "\n",
        "    for epoch in range(current_epoch + 1, current_epoch + more_epoch + 1):\n",
        "\n",
        "        current_epoch = epoch\n",
        "\n",
        "        for _ in range(epoch_scale):\n",
        "            loss = torch.tensor(0.).cuda()\n",
        "            inner_loss = torch.tensor(0.).cuda()\n",
        "            for _ in range(batch_scale):\n",
        "                input, part_label = next(train_iter)\n",
        "                part_label = part_label.cuda()\n",
        "                for iperm in get_trans(num_trans):\n",
        "                    \n",
        "                    features = linear(model(*input, perm=iperm))\n",
        "                    logits_all = model.part_classfier(features)\n",
        "                    loss = loss + crit(logits_all.view(-1, num_parts), part_label.view(-1))\n",
        "\n",
        "                    with torch.no_grad():\n",
        "                        accuracy.add((logits_all.max(dim=-1).indices == part_label).float().mean())\n",
        "                        miou.add(calc_miou(logits_all, part_label))\n",
        "\n",
        "                    if model_name == 'res_segment':\n",
        "                        inner_features = model.inner_ans\n",
        "                        logits_inner_all = model.inner_part_classfier(inner_features)\n",
        "                        inner_loss = inner_loss + crit(logits_inner_all.view(-1, num_parts), part_label.view(-1))\n",
        "          \n",
        "                    epoch_since += 1\n",
        "            \n",
        "            assert loss.isnan().sum() == 0\n",
        "            cum_loss += loss.item()\n",
        "            cum_inner_loss += inner_loss.item()\n",
        "            opt.zero_grad()\n",
        "            (loss + inner_loss).backward()\n",
        "            opt.step()\n",
        "            \n",
        "        \n",
        "\n",
        "        if cum_loss / epoch_since < threshold:\n",
        "            epoch_scale, batch_scale = batch_scale, epoch_scale\n",
        "            \n",
        "            logging.info(\"Threshold Reached\")\n",
        "            threshold = -1e10\n",
        "            \n",
        "        if epoch <= 5 or epoch % print_epoch == 0:\n",
        "            valid_str = \"\"\n",
        "            func = logging.debug\n",
        "\n",
        "            stop_training = False\n",
        "            if epoch == 5 or epoch % valid_epoch == 0:\n",
        "                vres, valid_str = evaluate(model, linear, valid_loader, noprint=True)\n",
        "                valid_str = \"valid = \" + valid_str\n",
        "                stop_training = (vres >= valid_result_threshold)\n",
        "                if vres > best_vres:\n",
        "                    best_vres = vres\n",
        "                    save(f\"best{prefix}\")\n",
        "                    valid_str += \" updated\"\n",
        "\n",
        "                func = logging.info\n",
        "            func(f\"train #{epoch} lr = {'%.2e' % sch.get_last_lr()[0]} loss = {'%.6lf / %.6lf' % (cum_loss / epoch_since, cum_inner_loss / epoch_since)} train = {accuracy} {miou} {valid_str}\")\n",
        "            epoch_since = cum_loss = cum_inner_loss = 0\n",
        "            accuracy.clear()\n",
        "            cloud_accuracy.clear()\n",
        "            miou.clear()\n",
        "\n",
        "            if stop_training:\n",
        "                break\n",
        "\n",
        "        sch.step()\n",
        "\n",
        "\n",
        "        if epoch % save_epoch == 0:\n",
        "            save(0)\n",
        "            tres, test_str = (best_vres, \"%.6lf\" % best_vres) if test_as_valid else evaluate(model, linear, test_loader, noprint=True, together=True)\n",
        "            logging.info(f\"Saved(Skipped) test = {test_str}\")\n",
        "\n",
        "        if epoch % cut_epoch == 0:\n",
        "            if batch_size > 8:\n",
        "                batch_size //= 2\n",
        "                epoch_scale *= 2\n",
        "\n",
        "            logging.info(f\"Cut batch_size = {batch_size} epoch_scale = {epoch_scale}\")\n"
      ],
      "metadata": {
        "id": "JWYGTAAvrbnq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Loop"
      ],
      "metadata": {
        "id": "Fx-MRVzNpY2B"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "train()"
      ],
      "metadata": {
        "id": "_fpT1UmBW5rz"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}