Run_PointTree_ShapeNetPart.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "[Release]Run_PointTree_ShapeNetPart.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "xJm_DG6kzGYG",
        "4lmXaxD5ynlo",
        "oLvduaqgSRuq",
        "-LzQmMFZmz60",
        "rqM9d9ravexM",
        "Y1XldfDwwWDO",
        "8V3dQ9Liwdey",
        "LaeKlPcjrdYV"
      ]
    },
    "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": "yMxwpnf886Im"
      },
      "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_shapenetpart'\n",
        "\n",
        "#encoder\n",
        "model_size = 2 ** 11\n",
        "sample_layers = 50 # 2\n",
        "channel = 1\n",
        "dim = 2048 # 2048\n",
        "dim_layer0 = 16 # 16\n",
        "dim_repeat_cut = 5\n",
        "use_sym = True\n",
        "\n",
        "#decoder \n",
        "ancestor_dim = 512 # 512\n",
        "model_name = \"segment\" #\"res_segment\"\n",
        "use_dyn_tree = False\n",
        "part_cls_dropout = None #0.5\n",
        "\n",
        "# data\n",
        "prefix = \"_affine_iter_xinf\"\n",
        "transform = affine_transform\n",
        "augment = 1\n",
        "no_prealign = False #True \n",
        "rotate_only = False #False\n",
        "load_pretrain = False\n",
        "augment_fn = lambda pts: augment_generator(pts) #, shift=True, scale=True, rotate_y_axis=True, agg_coef=0.5)\n",
        "pca_augment = False\n",
        "test_as_valid = False #False\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, shift=True, scale=True, rotate_y_axis=True, agg_coef=0.5)\n",
        "# pca_augment = False\n",
        "# test_as_valid = False #False\n",
        "# prefix = \"_homo_x30\"\n",
        "# transform = homo_transform\n",
        "# augment = 30\n",
        "# no_prealign = False\n",
        "# load_pretrain = False\n",
        "# pca_augment = True\n",
        "\n",
        "sample_child_first = False # True in l7s1\n",
        "num_classes = 16\n",
        "num_parts = 50\n",
        "DATASET = './datasets/ShapeNetPart'\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",
      "source": [
        "import json\n",
        "class_parts = json.load(open(f'{DATASET}/id2parts.json'))\n",
        "class_parts_list = []\n",
        "part_mapping = []\n",
        "for i, p in enumerate(class_parts):\n",
        "    class_parts_list.append(torch.tensor(p).cuda())\n",
        "    \n",
        "    mask = torch.zeros(num_parts).cuda().bool()\n",
        "    mask[p] = True\n",
        "    class_parts[i] = mask\n",
        "\n",
        "    pmap = torch.full([num_parts], -1).cuda()\n",
        "    pmap[p] = torch.arange(len(p)).cuda()\n",
        "    part_mapping.append(pmap)\n",
        "\n",
        "\n",
        "class_parts = torch.stack(class_parts, dim=0)\n"
      ],
      "metadata": {
        "id": "b-UHUjMLoBWk"
      },
      "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",
        "        encoder_fn = lambda point_dim : Encoder(model_size, sample_layers, dim, OUTPUT, point_dim=point_dim, channel=channel, sample_child_first=sample_child_first, dim_layer0=dim_layer0, dim_repeat_cut=dim_repeat_cut, use_symmetry_loss=use_sym).cuda()\n",
        "        model = ResSegment(ancestor_dim, encoder_fn, use_dyn_tree=use_dyn_tree, carry_dim_seg1=dim//16, num_parts=num_parts, part_cls_dropout=part_cls_dropout).cuda()\n",
        "    else:\n",
        "        encoder = Encoder(model_size, sample_layers, dim, OUTPUT, channel=channel, sample_child_first=sample_child_first, dim_layer0=dim_layer0, dim_repeat_cut=dim_repeat_cut, use_symmetry_loss=use_sym).cuda()\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"
      ],
      "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",
      "metadata": {
        "id": "vG7rbo4ekcjE"
      },
      "source": [
        "n_train = 6\n",
        "clouds = []\n",
        "labels = []\n",
        "extra_labels = []\n",
        "for i in range(n_train):\n",
        "    data_file = h5py.File(f'{DATASET}/train{i}.h5')\n",
        "    clouds.append(torch.tensor(np.array(data_file['data'])))\n",
        "    labels.append(torch.tensor(np.array(data_file['seg'])))\n",
        "    extra_labels.append(torch.tensor(np.array(data_file['label'])))\n",
        "\n",
        "clouds = torch.cat(clouds, dim=0)\n",
        "labels = torch.cat(labels, dim=0)\n",
        "extra_labels = torch.cat(extra_labels, dim=0)\n",
        "\n",
        "train_dataset = PointCloudDataset(clouds, labels, model.tree.arrange, augment=augment, transform=transform, make=make, extra_labels=extra_labels, augment_fn=augment_fn)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6C6m-AF1TMAA"
      },
      "source": [
        "n_valid = 1\n",
        "clouds = []\n",
        "labels = []\n",
        "extra_labels = []\n",
        "for i in range(n_valid):\n",
        "    data_file = h5py.File(f'{DATASET}/val{i}.h5')\n",
        "    clouds.append(torch.tensor(np.array(data_file['data'])))\n",
        "    labels.append(torch.tensor(np.array(data_file['seg'])))\n",
        "    extra_labels.append(torch.tensor(np.array(data_file['label'])))\n",
        "\n",
        "clouds = torch.cat(clouds, dim=0)\n",
        "labels = torch.cat(labels, dim=0)\n",
        "extra_labels = torch.cat(extra_labels, dim=0)\n",
        "\n",
        "valid_dataset = PointCloudDataset(clouds, labels, model.tree.arrange, augment=1, transform=transform, make=make, extra_labels=extra_labels)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "M3j6QUPo11OV"
      },
      "source": [
        "n_test = 2\n",
        "clouds = []\n",
        "labels = []\n",
        "extra_labels = []\n",
        "for i in range(n_test):\n",
        "    data_file = h5py.File(f'{DATASET}/test{i}.h5')\n",
        "    clouds.append(torch.tensor(np.array(data_file['data'])))\n",
        "    labels.append(torch.tensor(np.array(data_file['seg'])))\n",
        "    extra_labels.append(torch.tensor(np.array(data_file['label'])))\n",
        "\n",
        "clouds = torch.cat(clouds, dim=0)\n",
        "labels = torch.cat(labels, dim=0)\n",
        "extra_labels = torch.cat(extra_labels, dim=0)\n",
        "\n",
        "test_dataset = PointCloudDataset(clouds, labels, model.tree.arrange, augment=1, make=make, transform=transform, extra_labels=extra_labels)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "if test_as_valid:\n",
        "    a = train_dataset\n",
        "    b = valid_dataset\n",
        "    train_dataset = PointCloudDataset(torch.cat([a.clouds, b.clouds], dim=0), torch.cat([a.labels, b.labels], dim=0), model.tree.arrange, augment=1, transform=transform, make=make, extra_labels=torch.cat([a.extra_labels, b.extra_labels], dim=0))\n",
        "    valid_dataset = test_dataset\n"
      ],
      "metadata": {
        "id": "bLtag9_buIjd"
      },
      "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": "code",
      "metadata": {
        "id": "MfWwHySFzgbz"
      },
      "source": [
        "target = [] # ['train', 'valid', 'test']\n",
        "\n",
        "import gc\n",
        "gc.collect()\n",
        "\n",
        "for name, dataset in zip(['train', 'valid', 'test'], [train_dataset, valid_dataset, test_dataset]):\n",
        "    from tqdm import tqdm\n",
        "\n",
        "    if name not in target:\n",
        "        continue\n",
        "\n",
        "    num_workers = 16\n",
        "    batch_size = 8\n",
        "\n",
        "    data_init = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=placeholder, pin_memory=False, prefetch_factor=32, drop_last=False)\n",
        "    logging.info(f\"Init {name}\")\n",
        "\n",
        "    counter = 0\n",
        "    mem = []\n",
        "\n",
        "    for i, data in enumerate(tqdm(data_init)):\n",
        "        mem += data\n",
        "        # logging.debug(f\"Init {name}: {i+1}/{len(data_init)}\")\n",
        "\n",
        "        if len(mem) >= 16384:\n",
        "            torch.save(mem, f'{OUTPUT}/{name}_data{prefix}.{counter}.pth')\n",
        "            counter += 1\n",
        "            mem.clear()\n",
        "            gc.collect()\n",
        "            \n",
        "    torch.save(mem, f'{OUTPUT}/{name}_data{prefix}.{counter}.pth')\n",
        "    del mem\n",
        "    gc.collect()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oiILxo2qlfBf"
      },
      "source": [
        "# **Train**"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Defs"
      ],
      "metadata": {
        "id": "Y1XldfDwwWDO"
      }
    },
    {
      "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 = 128\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 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",
        "    # counted = in_label.sum(dim=-1) > 0.5\n",
        "    # counted = class_parts[cloud_label.long().cuda()]\n",
        "    \n",
        "    # appeared = in_label.sum(dim=-1) > 0.5\n",
        "    # if (counted != appeared).sum().item() != 0:\n",
        "    #     for cl, cnt, app in zip(cloud_label, counted, appeared):\n",
        "    #         if (cnt != app).sum().item() != 0:\n",
        "    #             arr = torch.arange(num_parts).cuda()\n",
        "    #             print(f\"Error cnt = {arr[cnt]} app = {arr[app]}\")\n",
        "    #     assert False\n",
        "\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], together=False):\n",
        "    cls = AccStat()\n",
        "    stat = AccStat()\n",
        "    trad = AccStat()\n",
        "\n",
        "    class_stat = [AccStat() for _ in range(num_classes)]\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, cloud_label) in enumerate(loader):\n",
        "        part_label = part_label.cuda()\n",
        "        cloud_label = cloud_label.squeeze(-1).cuda()\n",
        "        with torch.no_grad():\n",
        "            for iperm in perms:\n",
        "                if together:\n",
        "                    features = linear(model(*input, perm=iperm))\n",
        "                    logits_all = model.part_classfier(features)\n",
        "\n",
        "                    for cloud_l, part_l, logits in zip(cloud_label, part_label, logits_all):\n",
        "                        C = cloud_l.item()\n",
        "                        part_l = part_mapping[C][part_l]\n",
        "                        logits = logits[:, class_parts[C]]\n",
        "\n",
        "                        x = miou(logits, part_l)\n",
        "                        stat.add(x)\n",
        "                        class_stat[C].add(x)\n",
        "\n",
        "                else:\n",
        "                    features = linear(model(*input, perm=iperm))\n",
        "                    class_index = cloud_label.unique().item()\n",
        "\n",
        "                    logits = model.part_classfier(features)\n",
        "                    part_label = part_mapping[class_index][part_label]\n",
        "\n",
        "                    stat.add(miou(logits, part_label))\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",
        "    inst_miou = stat.result()\n",
        "    class_miou = sum([s.result() for s in class_stat]) / len(class_stat) if together else 0.0\n",
        "\n",
        "    if not noprint:\n",
        "        logging.info(\" \".join(map(str, class_stat)))\n",
        "    \n",
        "    return inst_miou, \"%.6lf %.6lf\" % (inst_miou, class_miou)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Train function"
      ],
      "metadata": {
        "id": "8V3dQ9Liwdey"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uBpg22qOl2Hu"
      },
      "source": [
        "import torch.nn.functional as F\n",
        "\n",
        "\n",
        "def train(class_index, more_epoch=10000, valid_result_threshold=999.0):\n",
        "    current_epoch = 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(len(class_parts_list[class_index]))\n",
        "    linear = nn.Identity()\n",
        "    model.train()\n",
        "    linear.train()\n",
        "    activate = lambda x : x\n",
        "\n",
        "    if load_pretrain:\n",
        "        ckpt = torch.load(f\"{OUTPUT}/trained_{class_index}_best_affined_pca_nosample_x10.pth\")\n",
        "        model.load_state_dict(ckpt['model'])\n",
        "        logging.info(\"Pretrain loaded\")\n",
        "\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",
        "    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_{class_index}_{epoch}.pth\")\n",
        "\n",
        "    global batch_size \n",
        "\n",
        "    subset_crit = lambda x : x[-1].item() == class_index\n",
        "    num_workers = 24\n",
        "    mbtrain = make_batch_generator(pca_augment=pca_augment and not no_prealign)\n",
        "    train_loader = torch.utils.data.DataLoader(Subset(train_dataset, subset_crit), 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(Subset(valid_dataset, subset_crit), 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(Subset(test_dataset, subset_crit), 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 class_index = {class_index} epoch = {current_epoch + 1} ~ {more_epoch} threshold = {valid_result_threshold} count = {len(train_loader) * batch_size} num_parts = {len(class_parts_list[class_index])}\")\n",
        "\n",
        "    cum_loss = 0\n",
        "    cum_inner_loss = 0\n",
        "\n",
        "    batch_scale = 1 # 1 if basic else 1\n",
        "    epoch_scale = 8\n",
        "    num_trans = 0 #0 if no_prealign else 1\n",
        "\n",
        "\n",
        "    print_epoch = 20\n",
        "    valid_epoch = 20\n",
        "    epoch_since = 0\n",
        "    save_epoch = 100\n",
        "    cut_epoch = 10000000000\n",
        "\n",
        "    cloud_accuracy = AccStat()\n",
        "    accuracy = AccStat()\n",
        "    trad_accuracy = AccStat()\n",
        "    cloud_classify_coef = 0.3\n",
        "    part_classify_coef = 0.2\n",
        "    class_spec_coef = 1\n",
        "\n",
        "    threshold = -1.0\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.9999)\n",
        "    save(current_epoch)\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, cloud_label = next(train_iter)\n",
        "                part_label = part_label.cuda()\n",
        "                cloud_label = cloud_label.squeeze(-1).cuda()\n",
        "                assert cloud_label.unique().item() == class_index\n",
        "                for iperm in get_trans(num_trans):\n",
        "                    \n",
        "                    features = linear(model(*input, perm=iperm))\n",
        "\n",
        "                    logits = model.part_classfier(features)\n",
        "                    part_label = part_mapping[class_index][part_label]\n",
        "\n",
        "                    # weight = get_parts_weight(part_label, num_parts=len(class_parts_list[class_index]))\n",
        "                    loss += F.cross_entropy(logits.transpose(-1, -2), part_label) * class_spec_coef\n",
        "\n",
        "                    if model_name == 'res_segment':\n",
        "                        inner_logits = model.inner_part_classfier(model.inner_ans)\n",
        "                        inner_loss += F.cross_entropy(inner_logits.transpose(-1, -2), part_label) * class_spec_coef\n",
        "\n",
        "                    with torch.no_grad():\n",
        "                        accuracy.add(miou(logits, part_label))\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",
        "        sch.step()\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 % 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} {valid_str}\")\n",
        "            epoch_since = cum_loss = cum_inner_loss = 0\n",
        "            accuracy.clear()\n",
        "            cloud_accuracy.clear()\n",
        "            trad_accuracy.clear()\n",
        "\n",
        "            if stop_training:\n",
        "                break\n",
        "\n",
        "\n",
        "        if epoch % save_epoch == 0:\n",
        "            if not test_as_valid:\n",
        "                save(epoch)\n",
        "            tres, test_str = (0., best_vres) if test_as_valid else evaluate(model, linear, test_loader, noprint=True)\n",
        "            logging.info(f\"Saved 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"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def train_together(more_epoch=100000, valid_result_threshold=999.0):\n",
        "    current_epoch = 0\n",
        "    best_vres = -1.0\n",
        "    \n",
        "    from random import choice, randint\n",
        "    import build_tree\n",
        "    import torch\n",
        "    \n",
        "    linear = nn.Identity()\n",
        "    model.train()\n",
        "    linear.train()\n",
        "    activate = lambda x : x\n",
        "\n",
        "    if load_pretrain:\n",
        "        ckpt = torch.load(f\"{OUTPUT}/trained_together_best_affined_pca_nosample_x10.pth\")\n",
        "        model.load_state_dict(ckpt['model'])\n",
        "        logging.info(\"Pretrain loaded\")\n",
        "\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 = 24\n",
        "    mbtrain = make_batch_generator(pca_augment=pca_augment and not no_prealign)\n",
        "    train_loader = torch.utils.data.DataLoader(BalanceDataset(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 = 4 # 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 = 20\n",
        "    epoch_since = 0\n",
        "    save_epoch = 100\n",
        "    cut_epoch = 10000000000\n",
        "\n",
        "    cloud_accuracy = AccStat()\n",
        "    accuracy = AccStat()\n",
        "    trad_accuracy = AccStat()\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",
        "\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_together_{epoch}.pth\")\n",
        "\n",
        "\n",
        "    save(current_epoch)\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, cloud_label = next(train_iter)\n",
        "                part_label = part_label.cuda()\n",
        "                cloud_label = cloud_label.squeeze(-1).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",
        "\n",
        "                    if model_name == 'res_segment':\n",
        "                        inner_features = model.inner_ans\n",
        "                        logits_inner_all = model.inner_part_classfier(inner_features)\n",
        "\n",
        "                        for cloud_l, part_l, logits, logits_inner in zip(cloud_label, part_label, logits_all, logits_inner_all):\n",
        "                            C = cloud_l.item()\n",
        "                            part_l = part_mapping[C][part_l]\n",
        "                            logits = logits[:, class_parts[C]]\n",
        "                            logits_inner = logits_inner[:, class_parts[C]]\n",
        "\n",
        "                            # count = (torch.arange(logits.shape[-1], device='cuda')[:, None] == part_l[None, :]).sum(dim=-1)\n",
        "                            # weight = 1.0 / count.clamp(min=1)\n",
        "\n",
        "                            loss += F.cross_entropy(logits / temperature, part_l) * class_spec_coef / batch_size\n",
        "                            inner_loss += F.cross_entropy(logits_inner / temperature, part_l)  * class_spec_coef / batch_size\n",
        "\n",
        "                            with torch.no_grad():\n",
        "                                accuracy.add(miou(logits, part_l))    \n",
        "\n",
        "\n",
        "                    else:\n",
        "                        assert False, \"disabled\"\n",
        "                        for cloud_l, part_l, logits in zip(cloud_label, part_label, logits_all):\n",
        "                            C = cloud_l.item()\n",
        "                            part_l = part_mapping[C][part_l]\n",
        "                            logits = logits[:, class_parts[C]]\n",
        "                            loss += F.cross_entropy(logits / temperature, part_l) * class_spec_coef / batch_size\n",
        "\n",
        "                            with torch.no_grad():\n",
        "                                accuracy.add(miou(logits, part_l))\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 % valid_epoch == 0:\n",
        "                vres, valid_str = evaluate(model, linear, valid_loader, noprint=True, together=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} {valid_str}\")\n",
        "            epoch_since = cum_loss = cum_inner_loss = 0\n",
        "            accuracy.clear()\n",
        "            cloud_accuracy.clear()\n",
        "            trad_accuracy.clear()\n",
        "\n",
        "            if stop_training:\n",
        "                break\n",
        "\n",
        "            if epoch % print_epoch == 0:\n",
        "                sch.step()\n",
        "\n",
        "\n",
        "        if epoch % save_epoch == 0:\n",
        "            if not test_as_valid:\n",
        "                save(epoch)\n",
        "            tres, test_str = (0., best_vres) if test_as_valid else evaluate(model, linear, test_loader, noprint=True, together=True)\n",
        "            logging.info(f\"Saved 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": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "a5W0z8dulaUT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Loops"
      ],
      "metadata": {
        "id": "0ridZ4M4wX6K"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "train(0, 2000)"
      ],
      "metadata": {
        "id": "oadU9hLAQ4d2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_together()"
      ],
      "metadata": {
        "id": "o1IyB2BdBkye"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}