modelnet40_train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "import h5py\n",
    "import sklearn.metrics as metrics\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch.nn.init as init\n",
    "import random\n",
    "\n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "\n",
    "class TransformNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(TransformNet, self).__init__()\n",
    "        self.k = 3\n",
    "        self.bn1 = nn.BatchNorm2d(64)\n",
    "        self.bn2 = nn.BatchNorm2d(128)\n",
    "        self.bn3 = nn.BatchNorm1d(1024)\n",
    "        self.conv1 = nn.Sequential(nn.Conv2d(6, 64, kernel_size=1, bias=False), self.bn1, nn.LeakyReLU(negative_slope=0.2))\n",
    "        self.conv2 = nn.Sequential(nn.Conv2d(64, 128, kernel_size=1, bias=False), self.bn2, nn.LeakyReLU(negative_slope=0.2))\n",
    "        self.conv3 = nn.Sequential(nn.Conv1d(128, 1024, kernel_size=1, bias=False), self.bn3, nn.LeakyReLU(negative_slope=0.2))\n",
    "        self.linear1 = nn.Linear(1024, 512, bias=False)\n",
    "        self.bn3 = nn.BatchNorm1d(512)\n",
    "        self.linear2 = nn.Linear(512, 256, bias=False)\n",
    "        self.bn4 = nn.BatchNorm1d(256)\n",
    "        self.transform = nn.Linear(256, 3 * 3)\n",
    "        init.constant_(self.transform.weight, 0)\n",
    "        init.eye_(self.transform.bias.view(3, 3))\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.size(0)\n",
    "        x = self.conv1(x)  # (batch_size, 3*2, num_points, k) -> (batch_size, 64, num_points, k)\n",
    "        x = self.conv2(x)  # (batch_size, 64, num_points, k) -> (batch_size, 128, num_points, k)\n",
    "        x = x.max(dim=-1, keepdim=False)[0]  # (batch_size, 128, num_points, k) -> (batch_size, 128, num_points)\n",
    "        x = self.conv3(x)  # (batch_size, 128, num_points) -> (batch_size, 1024, num_points)\n",
    "        x = x.max(dim=-1, keepdim=False)[0]  # (batch_size, 1024, num_points) -> (batch_size, 1024)\n",
    "        x = F.leaky_relu(self.bn3(self.linear1(x)), negative_slope=0.2)  # (batch_size, 1024) -> (batch_size, 512)\n",
    "        x = F.leaky_relu(self.bn4(self.linear2(x)), negative_slope=0.2)  # (batch_size, 512) -> (batch_size, 256)\n",
    "        x = self.transform(x)  # (batch_size, 256) -> (batch_size, 3*3)\n",
    "        x = x.view(batch_size, 3, 3)  # (batch_size, 3*3) -> (batch_size, 3, 3)\n",
    "        return x\n",
    "\n",
    "\n",
    "def knn(x, k):\n",
    "    inner = -2*torch.matmul(x.transpose(2, 1), x)\n",
    "    xx = torch.sum(x**2, dim=1, keepdim=True)\n",
    "    pairwise_distance = -xx - inner - xx.transpose(2, 1)\n",
    "    idx = pairwise_distance.topk(k=k, dim=-1)[1]   # (batch_size, num_points, k)\n",
    "    return idx\n",
    "\n",
    "\n",
    "def get_graph_feature(x, k=20, idx=None):\n",
    "    batch_size = x.size(0)\n",
    "    num_points = x.size(2)\n",
    "    x = x.view(batch_size, -1, num_points)\n",
    "    if idx is None:\n",
    "        idx = knn(x, k=k)   # (batch_size, num_points, k)\n",
    "    idx_base = torch.arange(0, batch_size, device=x.device).view(-1, 1, 1)*num_points\n",
    "    idx = idx + idx_base\n",
    "    idx = idx.view(-1)\n",
    "    _, num_dims, _ = x.size()\n",
    "    x = x.transpose(2, 1).contiguous()\n",
    "    feature = x.view(batch_size*num_points, -1)[idx, :]\n",
    "    feature = feature.view(batch_size, num_points, k, num_dims) \n",
    "    x = x.view(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1)\n",
    "    feature = torch.cat((feature-x, x), dim=3).permute(0, 3, 1, 2).contiguous()\n",
    "    return feature\n",
    "\n",
    "\n",
    "\n",
    "class DGCNN(nn.Module):\n",
    "    def __init__(self, k=25, output_channels=40):\n",
    "        super(DGCNN, self).__init__()\n",
    "        self.k = k\n",
    "        \n",
    "        self.bn1 = nn.BatchNorm2d(64)\n",
    "        self.bn2 = nn.BatchNorm2d(64)\n",
    "        self.bn3 = nn.BatchNorm2d(128)\n",
    "        self.bn4 = nn.BatchNorm2d(256)\n",
    "        self.bn5 = nn.BatchNorm1d(1024)\n",
    "        self.transform_net = TransformNet()\n",
    "\n",
    "        self.linearsel1 = nn.Conv1d(24, 512, kernel_size=1)\n",
    "        self.bnsel1 = nn.BatchNorm1d(512)\n",
    "        self.linearsel2 = nn.Conv1d(512, 256, kernel_size=1)\n",
    "        self.bnsel2 = nn.BatchNorm1d(256)\n",
    "        self.linearsel3 = nn.Linear(256, 24)\n",
    "\n",
    "        self.conv1 = nn.Sequential(nn.Conv2d(6, 64, kernel_size=1, bias=False), self.bn1, nn.LeakyReLU(negative_slope=0.2))\n",
    "        self.conv2 = nn.Sequential(nn.Conv2d(64*2, 64, kernel_size=1, bias=False), self.bn2, nn.LeakyReLU(negative_slope=0.2))\n",
    "        self.conv3 = nn.Sequential(nn.Conv2d(64*2, 128, kernel_size=1, bias=False), self.bn3, nn.LeakyReLU(negative_slope=0.2))\n",
    "        self.conv4 = nn.Sequential(nn.Conv2d(128*2, 256, kernel_size=1, bias=False), self.bn4, nn.LeakyReLU(negative_slope=0.2))\n",
    "        self.conv5 = nn.Sequential(nn.Conv1d(512, 1024, kernel_size=1, bias=False), self.bn5, nn.LeakyReLU(negative_slope=0.2))\n",
    "        \n",
    "        self.bn1 = nn.BatchNorm1d(64)\n",
    "        self.bn2 = nn.BatchNorm1d(64)\n",
    "        self.bn3 = nn.BatchNorm1d(128)\n",
    "        self.bn4 = nn.BatchNorm1d(256)\n",
    "        self.bn5 = nn.BatchNorm1d(1024)\n",
    "        \n",
    "        self.linear1 = nn.Linear(1024*2, 512, bias=False)\n",
    "        self.bn6 = nn.BatchNorm1d(512)\n",
    "        self.dp1 = nn.Dropout()\n",
    "        self.linear2 = nn.Linear(512, 512)\n",
    "        self.bn7 = nn.BatchNorm1d(512)\n",
    "        self.dp2 = nn.Dropout()\n",
    "        self.linear3 = nn.Linear(512, output_channels)\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        bs = x.size(0)\n",
    "        \n",
    "        # pose selector\n",
    "        xf = x.permute(0, 2, 3, 1).view(bs, 1024*3, 24)\n",
    "        s = F.leaky_relu(self.bnsel1(self.linearsel1(xf.transpose(2, 1))), negative_slope=0.2)\n",
    "        s = F.leaky_relu(self.bnsel2(self.linearsel2(s)), negative_slope=0.2)\n",
    "        s = F.adaptive_max_pool1d(s, 1).view(bs, -1)\n",
    "        s = F.softmax(self.linearsel3(s), dim=1).unsqueeze(-1)\n",
    "        x_int = torch.bmm(xf, s).view(bs, 1024, 3).permute(0 ,2, 1)  \n",
    "\n",
    "        x0 = get_graph_feature(x_int, k=self.k)  # (batch_size, 3, num_points) -> (batch_size, 3*2, num_points, k)\n",
    "        t = self.transform_net(x0)  # (batch_size, 3, 3)\n",
    "        x = x_int.transpose(2, 1)  # (batch_size, 3, num_points) -> (batch_size, num_points, 3)\n",
    "        x = torch.bmm(x, t)  # (batch_size, num_points, 3) * (batch_size, 3, 3) -> (batch_size, num_points, 3)\n",
    "        x = x.transpose(2, 1)\n",
    "        \n",
    "        x = get_graph_feature(x, k=self.k)\n",
    "        x = self.conv1(x)\n",
    "        x1 = x.max(dim=-1, keepdim=False)[0]\n",
    "        x = get_graph_feature(x1, k=self.k)\n",
    "        x = self.conv2(x)\n",
    "        x2 = x.max(dim=-1, keepdim=False)[0]\n",
    "        x = get_graph_feature(x2, k=self.k)\n",
    "        x = self.conv3(x)\n",
    "        x3 = x.max(dim=-1, keepdim=False)[0]\n",
    "        x = get_graph_feature(x3, k=self.k)\n",
    "        x = self.conv4(x)\n",
    "        x4 = x.max(dim=-1, keepdim=False)[0]\n",
    "        x = torch.cat((x1, x2, x3, x4), dim=1)\n",
    "        x = self.conv5(x)\n",
    "       \n",
    "        x1 = F.adaptive_max_pool1d(x, 1).view(bs, -1)\n",
    "        x2 = F.adaptive_avg_pool1d(x, 1).view(bs, -1)\n",
    "        \n",
    "        x = torch.cat((x1, x2), 1)\n",
    "        \n",
    "        x = F.leaky_relu(self.bn6(self.linear1(x)), negative_slope=0.2)\n",
    "        x = self.dp1(x)\n",
    "        x = F.leaky_relu(self.bn7(self.linear2(x)), negative_slope=0.2)\n",
    "        x = self.dp2(x)\n",
    "        x = self.linear3(x)\n",
    "        return x\n",
    "    \n",
    "    "
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "\n",
    "\n",
    "\n",
    "\"\"\"\n",
    "dataloader \n",
    "\"\"\"\n",
    "\n",
    "def generate_24_rotations():\n",
    "    res = []\n",
    "    for id in [[0, 1, 2], [1, 2, 0], [2, 0, 1]]:\n",
    "        R = np.identity(3)[:, id].astype(int)\n",
    "        R1= np.asarray([R[:, 0], R[:, 1], R[:, 2]]).T\n",
    "        R2 = np.asarray([-R[:, 0], -R[:, 1], R[:, 2]]).T\n",
    "        R3 = np.asarray([-R[:, 0], R[:, 1], -R[:, 2]]).T\n",
    "        R4 = np.asarray([R[:, 0], -R[:, 1], -R[:, 2]]).T\n",
    "        res += [R1, R2, R3, R4]\n",
    "    for id in [[0, 2, 1], [1, 0, 2], [2, 1, 0]]:\n",
    "        R = np.identity(3)[:, id].astype(int)\n",
    "        R1 = np.asarray([-R[:, 0], -R[:, 1], -R[:, 2]]).T\n",
    "        R2 = np.asarray([-R[:, 0], R[:, 1], R[:, 2]]).T\n",
    "        R3 = np.asarray([R[:, 0], -R[:, 1], R[:, 2]]).T\n",
    "        R4 = np.asarray([R[:, 0], R[:, 1], -R[:, 2]]).T\n",
    "        res += [R1, R2, R3, R4]\n",
    "    return res\n",
    "\n",
    "\n",
    "\n",
    "class H5Loader(Dataset):\n",
    "    def __init__(self, list_dir, file_list, partition, vcand):\n",
    "        self.data_list = []\n",
    "        self.partition = partition\n",
    "        self.all_R = generate_24_rotations()\n",
    "        self.vcand = vcand\n",
    "        for file_name in open(file_list):\n",
    "            self.data_list.append(os.path.join(list_dir, file_name).rstrip())\n",
    "\n",
    "    def __getitem__(self, ind):\n",
    "        file = h5py.File(self.data_list[ind], 'r', swmr=True)\n",
    "        data = file['data'][:]\n",
    "        pt = data[:1024, :].reshape(1, 1024, 3)\n",
    "        pclist = []\n",
    "        nums =np.arange(24)\n",
    "        #if training data, select pattern among 24 possible patterns\n",
    "        if self.partition == 'train':\n",
    "            nums = self.vcand[random.randint(0, 23)]\n",
    "        for i in range(24):\n",
    "            ptc = pt @ self.all_R[nums[i]]\n",
    "            pclist.append(ptc)\n",
    "            \n",
    "        data = np.concatenate(pclist)\n",
    "        data, label = torch.from_numpy(data), torch.from_numpy(file['label'][:])\n",
    "        file.close()\n",
    "        return data, label\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data_list)\n",
    "\n",
    "    "
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "\n",
    "\n",
    "\"\"\"\n",
    "training function\n",
    "\"\"\"\n",
    "\n",
    "def cal_loss(pred, label, smoothing_eps=0.2):\n",
    "    label = label.contiguous().view(-1)\n",
    "    n_class = pred.size(1)\n",
    "    one_hot = torch.zeros_like(pred).scatter(1, label.view(-1, 1), 1)\n",
    "    one_hot = one_hot * (1 - smoothing_eps) + (1 - one_hot) * smoothing_eps / (n_class - 1)\n",
    "    log_prb = F.log_softmax(pred, dim=1)\n",
    "    loss = -(one_hot * log_prb).sum(dim=1).mean()\n",
    "    return loss\n",
    "\n",
    "\n",
    "\n",
    "def train(data_dir, log_dir, device, n_epoch=500, lr=1e-1, batch_size=16, n_class=40, smoothing=True):\n",
    "    # dataloader\n",
    "    vcand = np.loadtxt('./all_id.txt').astype(np.int)\n",
    "    train_set = H5Loader(os.path.join(data_dir,'train'), os.path.join(data_dir, 'train/train_list.txt'),'train', vcand)\n",
    "    test_set = H5Loader(os.path.join(data_dir,'test'), os.path.join(data_dir, 'test/test_list.txt'),'test', vcand)\n",
    "    train_generator = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=2, drop_last=True)\n",
    "    test_generator = DataLoader(test_set, batch_size=batch_size, shuffle=True, num_workers=2, drop_last=False)\n",
    "    train_size, valid_size = train_set.__len__(), test_set.__len__()\n",
    "    print('training data batches: {}, validation data batches: {}'.format(train_size, valid_size))   \n",
    "    \n",
    "    # initialization\n",
    "    model = DGCNN(output_channels=n_class).to(device)\n",
    "    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)\n",
    "    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, n_epoch, eta_min=1e-6)\n",
    "    max_test_acc = 0.0\n",
    "    logger = open(log_dir, 'w').close()\n",
    "    print('init done')\n",
    "    \n",
    "    # training and validation\n",
    "    for epoch in range(n_epoch):  \n",
    "        # training phase\n",
    "        loss_val = 0.0\n",
    "        model.train() \n",
    "        correct_train = 0.0\n",
    "        for inputs, labels in tqdm(train_generator):\n",
    "            inputs, labels = inputs.to(device).float(), labels.to(device).long().flatten()\n",
    "            optimizer.zero_grad()\n",
    "            pred = model(inputs)\n",
    "            loss = cal_loss(pred, labels)\n",
    "            correct_train += (torch.max(pred.data, dim=1)[1] == labels).sum().item()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            loss_val += loss.item()\n",
    "        scheduler.step()\n",
    "        curr_train_acc = correct_train / train_size\n",
    "        print('Epoch {}: avg_loss is {}, curr_train_acc is {:.5%};'.format(epoch, loss_val, curr_train_acc))\n",
    "\n",
    "        # testing phase\n",
    "        if epoch >= 1000 or epoch % 100 == 0:\n",
    "            model.eval()\n",
    "            test_pred, test_true = [], []\n",
    "            with torch.no_grad():\n",
    "                for inputs, labels in tqdm(test_generator):\n",
    "                    ids = torch.from_numpy(vcand).to(device).long()\n",
    "                    inputs, labels = inputs.to(device).float(), labels.to(device).squeeze()\n",
    "                    tmp_pred = []\n",
    "                    for vw in range(24):\n",
    "                        tmp_pred.append(model(inputs[:,ids[vw]]).detach().cpu().numpy().reshape(1,-1,n_class))\n",
    "                    pred = np.sum(np.concatenate(tmp_pred), axis=0)\n",
    "                    pred = np.argmax(pred,axis=1)\n",
    "                    test_pred.append(pred)\n",
    "                    test_true.append(labels.detach().cpu().numpy())\n",
    "            \n",
    "                curr_test_acc = metrics.accuracy_score(np.concatenate(test_true), np.concatenate(test_pred))\n",
    "                if max_test_acc < curr_test_acc:\n",
    "                    max_test_acc = curr_test_acc\n",
    "                    torch.save(model.state_dict(), './modelnet40_checkpoint.t7')\n",
    "                print('Epoch {}: curr_test_acc is {:.2%};  max_test_acc is {:.2%}'.format(epoch, curr_test_acc, max_test_acc))\n",
    "          \n",
    "            # logger to txt\n",
    "            f = open(log_dir, 'a')\n",
    "            f.write('Epoch {}: loss {}, curr_acc_train {:.2%}, curr_acc {:.2%},  max_acc is {:.2%}\\n'.format(epoch, loss_val, curr_train_acc, curr_test_acc, max_test_acc))\n",
    "            f.close()\n",
    "   \n",
    "    print('training finished')\n",
    "\n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "\n",
    "if __name__ == '__main__':\n",
    "    \n",
    "    #directories need to be changed\n",
    "    data_dir = './dataset/modelnet40/pca/'\n",
    "    log_dir = 'modelnet40_log.txt'\n",
    "    device = torch.device('cuda:2')\n",
    "    \n",
    "    torch.manual_seed(123)\n",
    "    torch.cuda.manual_seed(123)\n",
    "    \n",
    "    \n",
    "    train(data_dir, log_dir, device=device, n_epoch=4000, lr=1e-1, batch_size=32, n_class=40, smoothing=True)\n",
    "\n"
   ],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}