SAFNet.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "SAFNet.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "RkUv6C2AaqOT"
      },
      "source": [
        "import numpy as np\n",
        "import scipy.io as sio\n",
        "from sklearn.decomposition import PCA\n",
        "from sklearn.model_selection import train_test_split\n",
        "import torch, math\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from functools import reduce\n",
        "import os\n",
        "import random\n",
        "from skimage import io, measure"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e_ySW9czhD5Y"
      },
      "source": [
        "Setting the seed of GPU"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5cuA86GuhESf"
      },
      "source": [
        "# def seed_torch(seed = 123):\n",
        "#     random.seed(seed)\n",
        "#     os.environ['PYTHONHASHSEED'] = str(seed) \n",
        "#     np.random.seed(seed)\n",
        "#     torch.manual_seed(seed)\n",
        "#     torch.cuda.manual_seed(seed)\n",
        "#     torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.\n",
        "#     torch.backends.cudnn.benchmark = False\n",
        "#     torch.backends.cudnn.deterministic = True\n",
        "\n",
        "# seed_torch()"
      ],
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jO_FCfwRm9am"
      },
      "source": [
        "# 1. Feature extraction network"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SVRw577qaqzU"
      },
      "source": [
        "class FeatNet(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(FeatNet, self).__init__()\n",
        "\n",
        "        self.conv1 = nn.Conv2d(1, 16, 3, 1, 1)\n",
        "        self.bn1_1 = nn.BatchNorm2d(16)\n",
        "        \n",
        "        self.conv1_1 = nn.Conv2d(16, 16, 3, 1, 1)\n",
        "        self.bn1_1 = nn.BatchNorm2d(16)\n",
        "        self.conv1_2 = nn.Conv2d(16, 16, 3, 1, 1)\n",
        "        self.bn1_2 = nn.BatchNorm2d(16)\n",
        "\n",
        "        self.conv2 = nn.Conv2d(16, 32, 1 ,2, 1)\n",
        "\n",
        "        self.conv2_1 = nn.Conv2d(32, 32, 3, 1, 1)\n",
        "        self.bn2_1 = nn.BatchNorm2d(32)\n",
        "        self.conv2_2 = nn.Conv2d(32, 32, 3, 1, 1)\n",
        "        self.bn2_2 = nn.BatchNorm2d(32)\n",
        "\n",
        "        self.conv3 = nn.Conv2d(32, 64, 1, 2, 1)\n",
        "\n",
        "        self.conv3_1 = nn.Conv2d(64, 64, 3, 1, 1)\n",
        "        self.bn3_1 = nn.BatchNorm2d(64)\n",
        "        self.conv3_2 = nn.Conv2d(64, 64, 3, 1, 1)\n",
        "        self.bn3_2 = nn.BatchNorm2d(64)\n",
        "\n",
        "        # Feature fusion\n",
        "        self.conv_fusion1 = nn.Conv2d(16, 64, 1, 4, 2)\n",
        "        self.conv_fusion2 = nn.Conv2d(32, 64, 1, 2, 1)\n",
        "\n",
        "    def forward(self, x):\n",
        "\n",
        "        x1 = self.conv1(x)\n",
        "        x = F.relu(self.bn1_1(self.conv1_1(x1)))\n",
        "        x_1 = F.relu(self.bn1_2(self.conv1_2(x)))\n",
        "        x = x1 + x_1\n",
        "        x2 = self.conv2(x)\n",
        "        x = F.relu(self.bn2_1(self.conv2_1(x2)))\n",
        "        x_2 = F.relu(self.bn2_2(self.conv2_2(x)))\n",
        "        x = x2 + x_2\n",
        "        x3 = self.conv3(x)\n",
        "        x = F.relu(self.bn3_1(self.conv3_1(x3)))\n",
        "        x_3 = F.relu(self.bn3_2(self.conv3_2(x)))\n",
        "        return x_1, x_2, x_3\n",
        "\n",
        "class FeatFuse(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(FeatFuse, self).__init__()\n",
        "\n",
        "        self.conv_fusion1 = nn.Conv2d(16, 64, 1, 4, 3)\n",
        "        self.conv_fusion2 = nn.Conv2d(32, 64, 1, 2, 1)\n",
        "\n",
        "        self.global_pool = nn.AdaptiveAvgPool2d(1)\n",
        "        self.fc1 = nn.Linear(64, 8)\n",
        "        self.fc2 = nn.Linear(8, 64*3)\n",
        "        \n",
        "        self.softmax = nn.Softmax(dim=1)\n",
        "\n",
        "    def forward(self, x1, x2, x3):\n",
        "        \n",
        "        batch_size = x1.size(0)\n",
        "        out_channels = x3.size(1)\n",
        "        x1 = self.conv_fusion1(x1)\n",
        "        x2 = self.conv_fusion2(x2)\n",
        "        output = []\n",
        "        output.append(x1)\n",
        "        output.append(x2)\n",
        "        output.append(x3)\n",
        "        x = x1 + x2 + x3\n",
        "          \n",
        "        x = self.global_pool(x)\n",
        "        x = x.view(x.size(0), -1)\n",
        "        x = F.relu(self.fc1(x))\n",
        "        x = self.fc2(x)\n",
        "        \n",
        "        a_b = x.reshape(batch_size, 3, out_channels, -1)\n",
        "        a_b = self.softmax(a_b)\n",
        "        #the part of selection\n",
        "        a_b = list(a_b.chunk(3, dim=1))#split to a and b\n",
        "        a_b = list(map(lambda x:x.reshape(batch_size, out_channels, 1, 1), a_b))\n",
        "        V = list(map(lambda x,y:x*y, output, a_b))\n",
        "        V = reduce(lambda x,y:x+y, V)\n",
        "        return V"
      ],
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_qIEQpWvnJj2"
      },
      "source": [
        "# 2. The proposed SAFNet"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-Hc5Au4Gbm0g"
      },
      "source": [
        "class Net(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Net, self).__init__()\n",
        "        \n",
        "        self.featnet = FeatNet()\n",
        "        self.featfuse = FeatFuse()\n",
        "        self.featnet1 = FeatNet()\n",
        "        self.featfuse1 = FeatFuse()\n",
        "        #self.global_pool = nn.AdaptiveAvgPool2d(1)\n",
        "        self.fc = nn.Linear(64, 2)\n",
        "        \n",
        "        self.global_pool1 = nn.AdaptiveAvgPool2d(1)\n",
        "        self.global_pool2 = nn.AdaptiveAvgPool2d(1)\n",
        "        self.fc1 = nn.Linear(64, 2)\n",
        "        self.fc2 = nn.Linear(64, 2)\n",
        "      \n",
        "\n",
        "    def forward(self, x, y):\n",
        "\n",
        "        x1_1, x1_2, x1_3 = self.featnet(x)\n",
        "        x2_1, x2_2, x2_3 = self.featnet1(y)\n",
        "\n",
        "        feat_11 = self.featfuse(x1_1, x1_2, x1_3)\n",
        "        feat_22 = self.featfuse1(x2_1, x2_2, x2_3)\n",
        "        feat_1 = self.global_pool1(feat_11)\n",
        "        feat_2 = self.global_pool2(feat_22)\n",
        "        feat_1 = feat_1.view(feat_1.size(0), -1)\n",
        "        feat_2 = feat_2.view(feat_2.size(0), -1)\n",
        "        feat_1 = self.fc1(feat_1)\n",
        "        feat_2 = self.fc2(feat_2)\n",
        "\n",
        "        feature_corr = self.xcorr_depthwise(feat_11, feat_22)\n",
        "        feat = feature_corr.view(feature_corr.size(0), -1)\n",
        "        #feat = global_pool(feature_corr)\n",
        "        feat = self.fc(feat)\n",
        "        return feat_1, feat_2, feat\n",
        "\n",
        "    def xcorr_depthwise(self, x, kernel):\n",
        "\n",
        "        batch = kernel.size(0)\n",
        "        channel = kernel.size(1)\n",
        "        x = x.view(1, batch*channel, x.size(2), x.size(3))\n",
        "        kernel = kernel.view(batch*channel, 1, kernel.size(2), kernel.size(3))\n",
        "        out = F.conv2d(x, kernel, groups=batch*channel)\n",
        "        out = out.view(batch, channel, out.size(2), out.size(3))\n",
        "        return out"
      ],
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iLsdHg0ZqjhO",
        "outputId": "47e2dcbc-7717-4173-f434-851ae8e23677"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ljW0mDornMLP"
      },
      "source": [
        "# 3. Data processing function"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NOnxQkLtbtq1"
      },
      "source": [
        "def addZeroPadding(X, margin=2):\n",
        "    newX = np.zeros((\n",
        "        X.shape[0] + 2 * margin,\n",
        "        X.shape[1] + 2 * margin,\n",
        "        X.shape[2]\n",
        "              ))\n",
        "    newX[margin:X.shape[0]+margin, margin:X.shape[1]+margin, :] = X\n",
        "    return newX\n",
        "\n",
        "def createImgCube(X ,gt ,pos:list ,windowSize=25):\n",
        "    margin = (windowSize-1)//2\n",
        "    zeroPaddingX = addZeroPadding(X, margin=margin)\n",
        "    dataPatches = np.zeros((pos.__len__(), windowSize, windowSize, X.shape[2]))\n",
        "    if( pos[-1][1]+1 != X.shape[1] ):\n",
        "        nextPos = (pos[-1][0] ,pos[-1][1]+1)\n",
        "    elif( pos[-1][0]+1 != X.shape[0] ):\n",
        "        nextPos = (pos[-1][0]+1 ,0)\n",
        "    else:\n",
        "        nextPos = (0,0)\n",
        "    return np.array([zeroPaddingX[i:i+windowSize, j:j+windowSize, :] for i,j in pos ]),\\\n",
        "    np.array([gt[i,j] for i,j in pos]) ,\\\n",
        "    nextPos\n",
        "\n",
        "def createPos(shape:tuple, pos:tuple, num:int):\n",
        "    if (pos[0]+1)*(pos[1]+1)+num >shape[0]*shape[1]:\n",
        "        num = shape[0]*shape[1]-( (pos[0])*shape[1] + pos[1] )\n",
        "    return [(pos[0]+(pos[1]+i)//shape[1] , (pos[1]+i)%shape[1] ) for i in range(num) ]\n",
        "\n",
        "def createPosWithoutZero(hsi, gt):\n",
        "\n",
        "    mask = gt > 0\n",
        "    return [(i,j) for i , row  in enumerate(mask) for j , row_element in enumerate(row) if row_element]\n",
        "\n",
        "def splitTrainTestSet(X, gt, testRatio, randomState=111):\n",
        "\n",
        "    X_train, X_test, gt_train, gt_test = train_test_split(X, gt, test_size=testRatio, random_state=randomState, stratify=gt)\n",
        "    return X_train, X_test, gt_train, gt_test\n",
        "\n",
        "def createImgPatch(lidar, pos:list, windowSize=25):\n",
        "\n",
        "    margin = (windowSize-1)//2\n",
        "    zeroPaddingLidar = np.zeros((\n",
        "      lidar.shape[0] + 2 * margin,\n",
        "      lidar.shape[1] + 2 * margin\n",
        "            ))\n",
        "    zeroPaddingLidar[margin:lidar.shape[0]+margin, margin:lidar.shape[1]+margin] = lidar\n",
        "    return np.array([zeroPaddingLidar[i:i+windowSize, j:j+windowSize] for i,j in pos ])\n",
        "\n",
        "def minmax_normalize(array):    \n",
        "    amin = np.min(array)\n",
        "    amax = np.max(array)\n",
        "    return (array - amin) / (amax - amin)\n",
        "def postprocess(res):\n",
        "    res_new = res\n",
        "    res = measure.label(res, connectivity=2)\n",
        "    num = res.max()\n",
        "    for i in range(1, num+1):\n",
        "        idy, idx = np.where(res==i)\n",
        "        if len(idy) <= 20:\n",
        "            res_new[idy, idx] = 0\n",
        "    return res_new"
      ],
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9Zwv_LyinO3t"
      },
      "source": [
        "# 4. Create dataloader"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fjjz-kyvbxDT"
      },
      "source": [
        "windowSize = 7 # patch size\n",
        "class_num = 2\n",
        "testRatio = 0.2 # the ratio of Validation set\n",
        "trainRatio = 0.9 # the ratio of Training set selected from preclassification\n",
        "\n",
        "data_path = '/content/drive/MyDrive/data'\n",
        "data_traingt = sio.loadmat(os.path.join(data_path, 'mask_train.mat'))['mask_train']\n",
        "data_testgt = sio.loadmat(os.path.join(data_path, 'mask_test.mat'))['mask_test']\n",
        "im1 = sio.loadmat(os.path.join(data_path, 'data_1.mat'))['data']\n",
        "im2 = sio.loadmat(os.path.join(data_path, 'data_2.mat'))['data']\n",
        "\n",
        "im1 = im1.reshape(im1.shape[0], im1.shape[1], 1)\n",
        "im2 = im2.reshape(im2.shape[0], im2.shape[1], 1)\n",
        "\n",
        "print (im1.shape)\n",
        "height , width, c = im1.shape\n",
        "\n",
        "# All pseudo-label set\n",
        "train_1, labels ,_ = createImgCube(im1, data_traingt, createPosWithoutZero(im1, data_traingt), windowSize=windowSize)\n",
        "train_2, _ ,_ = createImgCube(im2, data_traingt, createPosWithoutZero(im2, data_traingt), windowSize=windowSize)\n",
        "\n",
        "# training set selected from pseudo-label set\n",
        "train_1, _, train_labels, _ = splitTrainTestSet(train_1, labels, trainRatio, randomState=111)\n",
        "train_2, _, _, _ = splitTrainTestSet(train_2, labels, trainRatio, randomState=111)\n",
        "\n",
        "# data augmentation if need\n",
        "Xh = []\n",
        "Xl = []\n",
        "y = []\n",
        "for i in range(train_1.shape[0]):\n",
        "    Xh.append(train_1[i])\n",
        "    Xl.append(train_2[i])\n",
        "\n",
        "    noise = np.random.normal(0.0, 0.01, size=train_1[0].shape)\n",
        "    noise2 = np.random.normal(0.0, 0.01, size=train_2[0].shape)\n",
        "    Xh.append(np.flip(train_1[i] + noise, axis=1))\n",
        "    Xl.append(np.flip(train_2[i] + noise2, axis=1))\n",
        "\n",
        "    k = np.random.randint(4)\n",
        "    Xh.append(np.rot90(train_1[i], k=k))\n",
        "    Xl.append(np.rot90(train_2[i], k=k))\n",
        "\n",
        "    y.append(train_labels[i])\n",
        "    y.append(train_labels[i])\n",
        "    y.append(train_labels[i])\n",
        "\n",
        "labels = np.asarray(y, dtype=np.int8)\n",
        "train_1 = np.asarray(Xh, dtype=np.float32)\n",
        "train_2 = np.asarray(Xl,dtype=np.float32)\n",
        "train_1 = torch.from_numpy(train_1.transpose(0,3,1,2)).float()\n",
        "train_2 = torch.from_numpy(train_2.transpose(0,3,1,2)).float()\n",
        "\n",
        "# Select a partial validation set from the training set\n",
        "X_train, X_val, train_labels, val_labels = splitTrainTestSet(train_1, labels, testRatio, randomState=111)\n",
        "X_train_2, X_val_2, _, _ = splitTrainTestSet(train_2, labels, testRatio, randomState=111)\n",
        "\n",
        "# testing set\n",
        "X_test, test_labels ,_ = createImgCube(im1, data_traingt, createPosWithoutZero(im1, data_testgt), windowSize=windowSize)\n",
        "X_test_2, _ ,_ = createImgCube(im2, data_traingt, createPosWithoutZero(im2, data_testgt), windowSize=windowSize)\n",
        "X_test = torch.from_numpy(X_test.transpose(0,3,1,2)).float()\n",
        "X_test_2 = torch.from_numpy(X_test_2.transpose(0,3,1,2)).float()\n",
        "\n",
        "print (X_train.shape)\n",
        "print (X_val.shape)\n",
        "print (X_test.shape)\n",
        "print(\"Creating dataloader\")\n",
        "\n",
        "\"\"\" Training dataset\"\"\"\n",
        "class TrainDS(torch.utils.data.Dataset):\n",
        "    def __init__(self):\n",
        "        self.len = train_labels.shape[0]\n",
        "        self.hsi = torch.FloatTensor(X_train)\n",
        "        self.lidar = torch.FloatTensor(X_train_2)\n",
        "        self.labels = torch.LongTensor(train_labels - 1)\n",
        "    def __getitem__(self, index):\n",
        "        return self.hsi[index], self.lidar[index], self.labels[index]\n",
        "    def __len__(self):\n",
        "        return self.len\n",
        "\n",
        "\"\"\" Testing dataset\"\"\"\n",
        "class TestDS(torch.utils.data.Dataset):\n",
        "    def __init__(self):\n",
        "        self.len = test_labels.shape[0]\n",
        "        self.hsi = torch.FloatTensor(X_test)\n",
        "        self.lidar = torch.FloatTensor(X_test_2)\n",
        "        self.labels = torch.LongTensor(test_labels - 1)\n",
        "    def __getitem__(self, index):\n",
        "        return self.hsi[index], self.lidar[index], self.labels[index]\n",
        "    def __len__(self):\n",
        "        return self.len\n",
        "\"\"\" Validation dataset\"\"\"\t\t\n",
        "class ValDS(torch.utils.data.Dataset):\n",
        "    def __init__(self):\n",
        "        self.len = val_labels.shape[0]\n",
        "        self.hsi = torch.FloatTensor(X_val)\n",
        "        self.lidar = torch.FloatTensor(X_val_2)\n",
        "        self.labels = torch.LongTensor(val_labels - 1)\n",
        "    def __getitem__(self, index):\n",
        "        return self.hsi[index], self.lidar[index], self.labels[index]\n",
        "    def __len__(self):\n",
        "        return self.len\n",
        "\n",
        "\n",
        "# generate trainloader and valloader\n",
        "trainset = TrainDS() \n",
        "testset  = TestDS()\n",
        "valset = ValDS() \n",
        "train_loader = torch.utils.data.DataLoader(dataset = trainset, batch_size = 128, shuffle = True, num_workers = 0)\n",
        "test_loader = torch.utils.data.DataLoader(dataset = testset, batch_size =128, shuffle = False, num_workers = 0)\n",
        "val_loader = torch.utils.data.DataLoader(dataset = valset, batch_size =128, shuffle = False, num_workers = 0)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bQAMHkFinTkl"
      },
      "source": [
        "# 6. Running"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mnxnPMpxb0XA"
      },
      "source": [
        "class ContrastiveLoss(torch.nn.Module):\n",
        "    def __init__(self, margin=2.0):\n",
        "        super(ContrastiveLoss, self).__init__()\n",
        "        self.margin = margin\n",
        "\n",
        "    def forward(self, output1, output2, label):\n",
        "        euclidean_distance = F.pairwise_distance(output1, output2)\n",
        "        loss_contrastive = torch.mean((label) * torch.pow(euclidean_distance, 2)+(1-label) * torch.pow(torch.clamp(self.margin - euclidean_distance, min=0.0), 2))     \n",
        "        return loss_contrastive\n",
        "\n",
        "def calc_loss(x1, x2, outputs, labels, alpha):\n",
        "\n",
        "    criterion = nn.CrossEntropyLoss()\n",
        "    loss1 = criterion(outputs, labels)\n",
        "\n",
        "    contrastive = ContrastiveLoss()\n",
        "    loss2 = contrastive(x1, x2, labels)\n",
        "\n",
        "    loss_sum = loss1 + alpha* loss2\n",
        "    return loss_sum\n",
        "\n",
        "def train(model, device, train_loader, optimizer, epoch):\n",
        "    model.train()\n",
        "    total_loss = 0\n",
        "    for i, (inputs_1, inputs_2, labels) in enumerate(train_loader):\n",
        "\n",
        "        inputs_1, inputs_2 = inputs_1.to(device), inputs_2.to(device)\n",
        "        labels = labels.to(device)\n",
        "\n",
        "        optimizer.zero_grad() \n",
        "        feat_1, feat_2, outputs = model(inputs_1, inputs_2)\n",
        "        loss = calc_loss(feat_1, feat_2, outputs, labels, alpha = 1)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        total_loss += loss.item()\n",
        "\n",
        "    print('[Epoch: %d]   [loss avg: %.4f]   [current loss: %.4f]' %(epoch + 1, total_loss/(epoch+1), loss.item()))\n",
        "\n",
        "def test(model, device, test_loader):\n",
        "    model.eval()\n",
        "    count = 0\n",
        "    for inputs_1, inputs_2, labels in test_loader:\n",
        "\n",
        "        inputs_1, inputs_2 = inputs_1.to(device), inputs_2.to(device)\n",
        "        _, _, outputs = model(inputs_1, inputs_2)\n",
        "        outputs = np.argmax(outputs.detach().cpu().numpy(), axis=1)\n",
        "\n",
        "        if count == 0:\n",
        "            y_pred_test =  outputs\n",
        "            test_labels = labels\n",
        "            count = 1\n",
        "        else:\n",
        "            y_pred_test = np.concatenate( (y_pred_test, outputs) )\n",
        "            test_labels = np.concatenate( (test_labels, labels) )\n",
        "    a = 0\n",
        "    for c in range(len(y_pred_test)):\n",
        "        if test_labels[c]==y_pred_test[c]:\n",
        "            a = a+1\n",
        "    acc = a/len(y_pred_test)*100\n",
        "    print('%.2f' %(a/len(y_pred_test)*100))\n",
        "\n",
        "    return acc"
      ],
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9AJM099jcAUj"
      },
      "source": [
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "model = Net().to(device)\n",
        "\n",
        "num_epochs = 50\n",
        "lr = 0.001\n",
        "momentum = 0.9\n",
        "betas = (0.9, 0.999)\n",
        "\n",
        "params_to_update = list(model.parameters())\n",
        "\n",
        "# optimizer = torch.optim.Adam(params_to_update, lr=lr, betas=betas)\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum,weight_decay=0.0005)\n",
        "\n",
        "best_acc = 0\n",
        "for epoch in range(num_epochs):\n",
        "    train(model, device, train_loader, optimizer, epoch)\n",
        "    acc = test(model, device, val_loader)\n",
        "    if acc >= best_acc:\n",
        "        best_acc = acc\n",
        "        print(\"Save model!\")\n",
        "        torch.save(model.state_dict(),'/content/drive/MyDrive/model/model.pth')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kbpq604ro7hW"
      },
      "source": [
        "# 7. Record the final change map"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mU6Aj1eilG0D"
      },
      "source": [
        "model = Net().eval().to(device)\n",
        "model.load_state_dict(torch.load('/content/drive/MyDrive/model/model.pth'))\n",
        "acc = test(model, device, test_loader)\n",
        "print ('The final accuracy is ', acc)\n",
        "margin = (windowSize-1)//2\n",
        "im1 = addZeroPadding(im1, margin=margin)\n",
        "im2 = addZeroPadding(im2, margin=margin)\n",
        "\n",
        "outputs = np.zeros((height, width))\n",
        "for i in range(height):\n",
        "    for j in range(width):\n",
        "    # if preclassify_lab[i, j]!= 1.5:\n",
        "    #   outputs[i, j] = preclassify_lab[i, j]\n",
        "\n",
        "    # else:\n",
        "        patch1 = im1[i:i+windowSize, j:j+windowSize, :]\n",
        "        patch1 = patch1.reshape(1, patch1.shape[0], patch1.shape[1], patch1.shape[2])\n",
        "        X_test_image = torch.FloatTensor(patch1.transpose(0, 3, 1, 2)).to(device)\n",
        "\n",
        "        patch2 = im2[i:i+windowSize, j:j+windowSize, :]\n",
        "        patch2 = patch2.reshape(1, patch2.shape[0], patch2.shape[1], patch2.shape[2])\n",
        "        X_test_image1 = torch.FloatTensor(patch2.transpose(0, 3, 1, 2)).to(device)\n",
        "\n",
        "        _, _, prediction = model(X_test_image, X_test_image1)\n",
        "        prediction = np.argmax(prediction.detach().cpu().numpy(), axis=1)\n",
        "        outputs[i][j] = prediction\n",
        "    if i % 20 == 0:\n",
        "        print('... ... row ', i, ' handling ... ...')\n",
        "# postprocessing if need\n",
        "outputs = postprocess(outputs)\n",
        "\n",
        "sio.savemat('result.mat', {'output': outputs})\n",
        "print('ALL Finish!!')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}