{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Introduction\n\nThis notebook shows how to train PraNet with TPUs on HPUMap Dataset.\n\n\n> we propose a parallel reverse attention network (PraNet) for accurate polyp segmentation in colonoscopy images. Specifically, we first aggregate ?the features in high-level layers using a parallel partial decoder (PPD). Based on the combined feature, we then generate a global map as the initial guidance area for the following components. In addition, we mine the boundary cues using the reverse attention (RA) module, which is able to establish the relationship between areas and boundary cues. Thanks to the recurrent cooperation mechanism between areas and boundaries, our PraNet is capable of calibrating some misaligned predictions, improving the segmentation accuracy. - Paper.\n\n![](https://github.com/DengPingFan/PraNet/raw/master/imgs/framework-final-min.png)"},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport numpy as np\nfrom PIL import Image\nimport torch.utils.data as data\nimport torchvision.transforms as transforms\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\nimport cv2 \n\nclass PolypDataset(data.Dataset):\n    \"\"\"\n    dataloader for polyp segmentation tasks\n    \"\"\"\n    def __init__(self, image_root, gt_root, trainsize):\n        self.trainsize = trainsize\n        \n        self.images = [image_root + f for f in os.listdir(image_root) if f.endswith('.jpg') or f.endswith('.png')]\n        self.gts = [gt_root + f for f in os.listdir(gt_root) if f.endswith('.jpg') or f.endswith('.png')]\n        self.images = sorted(self.images)\n        self.gts = sorted(self.gts)\n#         self.filter_files()\n        self.size = len(self.images)\n        self.img_transform =  A.Compose([\n                A.HorizontalFlip(),\n                A.VerticalFlip(),\n                A.RandomRotate90(),\n                A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=15, p=0.9, \n                                 border_mode=cv2.BORDER_REFLECT),\n                A.OneOf([\n                    A.OpticalDistortion(p=0.3),\n                    A.GridDistortion(p=.1),\n                    A.IAAPiecewiseAffine(p=0.3),\n                ], p=0.3),\n                A.OneOf([\n                    A.HueSaturationValue(10,15,10),\n                    A.CLAHE(clip_limit=2),\n                    A.RandomBrightnessContrast(),            \n                ], p=0.3),\n            A.Normalize(),\n            ToTensorV2()\n            ], p=1.0)\n\n    def __getitem__(self, index):\n        image = self.rgb_loader(self.images[index])\n        gt = self.binary_loader(self.gts[index])\n        image, gt =  self.resize(image, gt)\n        if self.img_transform:\n            # Applying augmentations if any. \n            sample = self.img_transform(image = image, \n                                     mask = gt)\n            image, gt = torch.reshape(sample['image'], (3, 512, 512)), torch.reshape(sample['mask'], (1, 512, 512))\n        return image, gt\n\n    def filter_files(self):\n        assert len(self.images) == len(self.gts)\n        images = []\n        gts = []\n        for img_path, gt_path in zip(self.images, self.gts):\n            img = Image.open(img_path)\n            gt = Image.open(gt_path)\n            if img.size == gt.size:\n                images.append(img_path)\n                gts.append(gt_path)\n        self.images = images\n        self.gts = gts\n\n    def rgb_loader(self, path):\n        with open(path, 'rb') as f:\n            img = Image.open(f)\n            return img.convert('RGB')\n        \n    def binary_loader(self, path):\n        with open(path, 'rb') as f:\n            img = Image.open(f)\n            # return img.convert('1')\n            return img.convert('L')\n\n    def resize(self, img, gt):\n        assert img.size == gt.size\n        h = self.trainsize\n        w = self.trainsize\n        return np.asarray(img.resize((w, h), Image.BILINEAR)), np.asarray(gt.resize((w, h), Image.NEAREST))\n\n\n    def __len__(self):\n        return self.size\n\n\ndef get_loader(image_root, gt_root, batchsize, trainsize, shuffle=True, num_workers=4, pin_memory=True):\n\n    dataset = PolypDataset(image_root, gt_root, trainsize)\n    data_loader = data.DataLoader(dataset=dataset,\n                                  batch_size=batchsize,\n                                  shuffle=shuffle,\n                                  num_workers=num_workers,\n                                  pin_memory=pin_memory)\n    return data_loader\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_root = '../input/hubmap-256x256/train/'\ngt_root = '../input/hubmap-256x256/masks/'\nbatchsize = 10\ntrainsize = 512\ntrain_loader = get_loader(image_root, gt_root, batchsize, trainsize, shuffle=True, num_workers=0, pin_memory=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch.nn as nn\nimport math\nimport torch.utils.model_zoo as model_zoo\nimport torch\nimport torch.nn.functional as F\n\n__all__ = ['Res2Net', 'res2net50_v1b', 'res2net101_v1b', 'res2net50_v1b_26w_4s']\n\nmodel_urls = {\n    'res2net50_v1b_26w_4s': 'https://shanghuagao.oss-cn-beijing.aliyuncs.com/res2net/res2net50_v1b_26w_4s-3cf99910.pth',\n    'res2net101_v1b_26w_4s': 'https://shanghuagao.oss-cn-beijing.aliyuncs.com/res2net/res2net101_v1b_26w_4s-0812c246.pth',\n}\n\n\nclass Bottle2neck(nn.Module):\n    expansion = 4\n\n    def __init__(self, inplanes, planes, stride=1, downsample=None, baseWidth=26, scale=4, stype='normal'):\n        \"\"\" Constructor\n        Args:\n            inplanes: input channel dimensionality\n            planes: output channel dimensionality\n            stride: conv stride. Replaces pooling layer.\n            downsample: None when stride = 1\n            baseWidth: basic width of conv3x3\n            scale: number of scale.\n            type: 'normal': normal set. 'stage': first block of a new stage.\n        \"\"\"\n        super(Bottle2neck, self).__init__()\n\n        width = int(math.floor(planes * (baseWidth / 64.0)))\n        self.conv1 = nn.Conv2d(inplanes, width * scale, kernel_size=1, bias=False)\n        self.bn1 = nn.BatchNorm2d(width * scale)\n\n        if scale == 1:\n            self.nums = 1\n        else:\n            self.nums = scale - 1\n        if stype == 'stage':\n            self.pool = nn.AvgPool2d(kernel_size=3, stride=stride, padding=1)\n        convs = []\n        bns = []\n        for i in range(self.nums):\n            convs.append(nn.Conv2d(width, width, kernel_size=3, stride=stride, padding=1, bias=False))\n            bns.append(nn.BatchNorm2d(width))\n        self.convs = nn.ModuleList(convs)\n        self.bns = nn.ModuleList(bns)\n\n        self.conv3 = nn.Conv2d(width * scale, planes * self.expansion, kernel_size=1, bias=False)\n        self.bn3 = nn.BatchNorm2d(planes * self.expansion)\n\n        self.relu = nn.ReLU(inplace=True)\n        self.downsample = downsample\n        self.stype = stype\n        self.scale = scale\n        self.width = width\n\n    def forward(self, x):\n        residual = x\n\n        out = self.conv1(x)\n        out = self.bn1(out)\n        out = self.relu(out)\n\n        spx = torch.split(out, self.width, 1)\n        for i in range(self.nums):\n            if i == 0 or self.stype == 'stage':\n                sp = spx[i]\n            else:\n                sp = sp + spx[i]\n            sp = self.convs[i](sp)\n            sp = self.relu(self.bns[i](sp))\n            if i == 0:\n                out = sp\n            else:\n                out = torch.cat((out, sp), 1)\n        if self.scale != 1 and self.stype == 'normal':\n            out = torch.cat((out, spx[self.nums]), 1)\n        elif self.scale != 1 and self.stype == 'stage':\n            out = torch.cat((out, self.pool(spx[self.nums])), 1)\n\n        out = self.conv3(out)\n        out = self.bn3(out)\n\n        if self.downsample is not None:\n            residual = self.downsample(x)\n\n        out += residual\n        out = self.relu(out)\n\n        return out\n\n\nclass Res2Net(nn.Module):\n\n    def __init__(self, block, layers, baseWidth=26, scale=4, num_classes=1000):\n        self.inplanes = 64\n        super(Res2Net, self).__init__()\n        self.baseWidth = baseWidth\n        self.scale = scale\n        self.conv1 = nn.Sequential(\n            nn.Conv2d(3, 32, 3, 2, 1, bias=False),\n            nn.BatchNorm2d(32),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(32, 32, 3, 1, 1, bias=False),\n            nn.BatchNorm2d(32),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(32, 64, 3, 1, 1, bias=False)\n        )\n        self.bn1 = nn.BatchNorm2d(64)\n        self.relu = nn.ReLU()\n        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n        self.layer1 = self._make_layer(block, 64, layers[0])\n        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)\n        self.avgpool = nn.AdaptiveAvgPool2d(1)\n        self.fc = nn.Linear(512 * block.expansion, num_classes)\n\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n            elif isinstance(m, nn.BatchNorm2d):\n                nn.init.constant_(m.weight, 1)\n                nn.init.constant_(m.bias, 0)\n\n    def _make_layer(self, block, planes, blocks, stride=1):\n        downsample = None\n        if stride != 1 or self.inplanes != planes * block.expansion:\n            downsample = nn.Sequential(\n                nn.AvgPool2d(kernel_size=stride, stride=stride,\n                             ceil_mode=True, count_include_pad=False),\n                nn.Conv2d(self.inplanes, planes * block.expansion,\n                          kernel_size=1, stride=1, bias=False),\n                nn.BatchNorm2d(planes * block.expansion),\n            )\n\n        layers = []\n        layers.append(block(self.inplanes, planes, stride, downsample=downsample,\n                            stype='stage', baseWidth=self.baseWidth, scale=self.scale))\n        self.inplanes = planes * block.expansion\n        for i in range(1, blocks):\n            layers.append(block(self.inplanes, planes, baseWidth=self.baseWidth, scale=self.scale))\n\n        return nn.Sequential(*layers)\n\n    def forward(self, x):\n        x = self.conv1(x)\n        x = self.bn1(x)\n        x = self.relu(x)\n        x = self.maxpool(x)\n\n        x = self.layer1(x)\n        x = self.layer2(x)\n        x = self.layer3(x)\n        x = self.layer4(x)\n\n        x = self.avgpool(x)\n        x = x.view(x.size(0), -1)\n        x = self.fc(x)\n\n        return x\n\n\ndef res2net50_v1b_26w_4s(pretrained=False, **kwargs):\n    \"\"\"Constructs a Res2Net-50_v1b_26w_4s lib.\n    Args:\n        pretrained (bool): If True, returns a lib pre-trained on ImageNet\n    \"\"\"\n    model = Res2Net(Bottle2neck, [3, 4, 6, 3], baseWidth=26, scale=4, **kwargs)\n    if pretrained:\n#         model_state = torch.load('/media/nercms/NERCMS/GepengJi/Medical_Seqmentation/CRANet/models/res2net50_v1b_26w_4s-3cf99910.pth')\n#         model.load_state_dict(model_state)\n        model.load_state_dict(model_zoo.load_url(model_urls['res2net50_v1b_26w_4s'], map_location = 'cpu'))\n    return model\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass BasicConv2d(nn.Module):\n    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1):\n        super(BasicConv2d, self).__init__()\n        self.conv = nn.Conv2d(in_planes, out_planes,\n                              kernel_size=kernel_size, stride=stride,\n                              padding=padding, dilation=dilation, bias=False)\n        self.bn = nn.BatchNorm2d(out_planes)\n        self.relu = nn.ReLU(inplace=True)\n\n    def forward(self, x):\n        x = self.conv(x)\n        x = self.bn(x)\n        return x\n\n\nclass RFB_modified(nn.Module):\n    def __init__(self, in_channel, out_channel):\n        super(RFB_modified, self).__init__()\n        self.relu = nn.ReLU(True)\n        self.branch0 = nn.Sequential(\n            BasicConv2d(in_channel, out_channel, 1),\n        )\n        self.branch1 = nn.Sequential(\n            BasicConv2d(in_channel, out_channel, 1),\n            BasicConv2d(out_channel, out_channel, kernel_size=(1, 3), padding=(0, 1)),\n            BasicConv2d(out_channel, out_channel, kernel_size=(3, 1), padding=(1, 0)),\n            BasicConv2d(out_channel, out_channel, 3, padding=3, dilation=3)\n        )\n        self.branch2 = nn.Sequential(\n            BasicConv2d(in_channel, out_channel, 1),\n            BasicConv2d(out_channel, out_channel, kernel_size=(1, 5), padding=(0, 2)),\n            BasicConv2d(out_channel, out_channel, kernel_size=(5, 1), padding=(2, 0)),\n            BasicConv2d(out_channel, out_channel, 3, padding=5, dilation=5)\n        )\n        self.branch3 = nn.Sequential(\n            BasicConv2d(in_channel, out_channel, 1),\n            BasicConv2d(out_channel, out_channel, kernel_size=(1, 7), padding=(0, 3)),\n            BasicConv2d(out_channel, out_channel, kernel_size=(7, 1), padding=(3, 0)),\n            BasicConv2d(out_channel, out_channel, 3, padding=7, dilation=7)\n        )\n        self.conv_cat = BasicConv2d(4*out_channel, out_channel, 3, padding=1)\n        self.conv_res = BasicConv2d(in_channel, out_channel, 1)\n\n    def forward(self, x):\n        x0 = self.branch0(x)\n        x1 = self.branch1(x)\n        x2 = self.branch2(x)\n        x3 = self.branch3(x)\n        x_cat = self.conv_cat(torch.cat((x0, x1, x2, x3), 1))\n\n        x = self.relu(x_cat + self.conv_res(x))\n        return x\n\n\nclass aggregation(nn.Module):\n    # dense aggregation, it can be replaced by other aggregation previous, such as DSS, amulet, and so on.\n    # used after MSF\n    def __init__(self, channel):\n        super(aggregation, self).__init__()\n        self.relu = nn.ReLU(True)\n\n        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n        self.conv_upsample1 = BasicConv2d(channel, channel, 3, padding=1)\n        self.conv_upsample2 = BasicConv2d(channel, channel, 3, padding=1)\n        self.conv_upsample3 = BasicConv2d(channel, channel, 3, padding=1)\n        self.conv_upsample4 = BasicConv2d(channel, channel, 3, padding=1)\n        self.conv_upsample5 = BasicConv2d(2*channel, 2*channel, 3, padding=1)\n\n        self.conv_concat2 = BasicConv2d(2*channel, 2*channel, 3, padding=1)\n        self.conv_concat3 = BasicConv2d(3*channel, 3*channel, 3, padding=1)\n        self.conv4 = BasicConv2d(3*channel, 3*channel, 3, padding=1)\n        self.conv5 = nn.Conv2d(3*channel, 1, 1)\n\n    def forward(self, x1, x2, x3):\n        x1_1 = x1\n        x2_1 = self.conv_upsample1(self.upsample(x1)) * x2\n        x3_1 = self.conv_upsample2(self.upsample(self.upsample(x1))) \\\n               * self.conv_upsample3(self.upsample(x2)) * x3\n\n        x2_2 = torch.cat((x2_1, self.conv_upsample4(self.upsample(x1_1))), 1)\n        x2_2 = self.conv_concat2(x2_2)\n\n        x3_2 = torch.cat((x3_1, self.conv_upsample5(self.upsample(x2_2))), 1)\n        x3_2 = self.conv_concat3(x3_2)\n\n        x = self.conv4(x3_2)\n        x = self.conv5(x)\n\n        return x\n\n\nclass PraNet(nn.Module):\n    # res2net based encoder decoder\n    def __init__(self, channel=32):\n        super(PraNet, self).__init__()\n        # ---- ResNet Backbone ----\n        self.resnet = res2net50_v1b_26w_4s(pretrained=True)\n        # ---- Receptive Field Block like module ----\n        self.rfb2_1 = RFB_modified(512, channel)\n        self.rfb3_1 = RFB_modified(1024, channel)\n        self.rfb4_1 = RFB_modified(2048, channel)\n        # ---- Partial Decoder ----\n        self.agg1 = aggregation(channel)\n        # ---- reverse attention branch 4 ----\n        self.ra4_conv1 = BasicConv2d(2048, 256, kernel_size=1)\n        self.ra4_conv2 = BasicConv2d(256, 256, kernel_size=5, padding=2)\n        self.ra4_conv3 = BasicConv2d(256, 256, kernel_size=5, padding=2)\n        self.ra4_conv4 = BasicConv2d(256, 256, kernel_size=5, padding=2)\n        self.ra4_conv5 = BasicConv2d(256, 1, kernel_size=1)\n        # ---- reverse attention branch 3 ----\n        self.ra3_conv1 = BasicConv2d(1024, 64, kernel_size=1)\n        self.ra3_conv2 = BasicConv2d(64, 64, kernel_size=3, padding=1)\n        self.ra3_conv3 = BasicConv2d(64, 64, kernel_size=3, padding=1)\n        self.ra3_conv4 = BasicConv2d(64, 1, kernel_size=3, padding=1)\n        # ---- reverse attention branch 2 ----\n        self.ra2_conv1 = BasicConv2d(512, 64, kernel_size=1)\n        self.ra2_conv2 = BasicConv2d(64, 64, kernel_size=3, padding=1)\n        self.ra2_conv3 = BasicConv2d(64, 64, kernel_size=3, padding=1)\n        self.ra2_conv4 = BasicConv2d(64, 1, kernel_size=3, padding=1)\n\n    def forward(self, x):\n        x = self.resnet.conv1(x)\n        x = self.resnet.bn1(x)\n        x = self.resnet.relu(x)\n        x = self.resnet.maxpool(x)      # bs, 64, 88, 88\n        # ---- low-level features ----\n        x1 = self.resnet.layer1(x)      # bs, 256, 88, 88\n        x2 = self.resnet.layer2(x1)     # bs, 512, 44, 44\n\n        x3 = self.resnet.layer3(x2)     # bs, 1024, 22, 22\n        x4 = self.resnet.layer4(x3)     # bs, 2048, 11, 11\n        x2_rfb = self.rfb2_1(x2)        # channel -> 32\n        x3_rfb = self.rfb3_1(x3)        # channel -> 32\n        x4_rfb = self.rfb4_1(x4)        # channel -> 32\n\n        ra5_feat = self.agg1(x4_rfb, x3_rfb, x2_rfb)\n        lateral_map_5 = F.interpolate(ra5_feat, scale_factor=8, mode='bilinear')    # NOTES: Sup-1 (bs, 1, 44, 44) -> (bs, 1, 352, 352)\n\n        # ---- reverse attention branch_4 ----\n        crop_4 = F.interpolate(ra5_feat, scale_factor=0.25, mode='bilinear')\n        x = -1*(torch.sigmoid(crop_4)) + 1\n        x = x.expand(-1, 2048, -1, -1).mul(x4)\n        x = self.ra4_conv1(x)\n        x = F.relu(self.ra4_conv2(x))\n        x = F.relu(self.ra4_conv3(x))\n        x = F.relu(self.ra4_conv4(x))\n        ra4_feat = self.ra4_conv5(x)\n        x = ra4_feat + crop_4\n        lateral_map_4 = F.interpolate(x, scale_factor=32, mode='bilinear')  # NOTES: Sup-2 (bs, 1, 11, 11) -> (bs, 1, 352, 352)\n\n        # ---- reverse attention branch_3 ----\n        crop_3 = F.interpolate(x, scale_factor=2, mode='bilinear')\n        x = -1*(torch.sigmoid(crop_3)) + 1\n        x = x.expand(-1, 1024, -1, -1).mul(x3)\n        x = self.ra3_conv1(x)\n        x = F.relu(self.ra3_conv2(x))\n        x = F.relu(self.ra3_conv3(x))\n        ra3_feat = self.ra3_conv4(x)\n        x = ra3_feat + crop_3\n        lateral_map_3 = F.interpolate(x, scale_factor=16, mode='bilinear')  # NOTES: Sup-3 (bs, 1, 22, 22) -> (bs, 1, 352, 352)\n\n        # ---- reverse attention branch_2 ----\n        crop_2 = F.interpolate(x, scale_factor=2, mode='bilinear')\n        x = -1*(torch.sigmoid(crop_2)) + 1\n        x = x.expand(-1, 512, -1, -1).mul(x2)\n        x = self.ra2_conv1(x)\n        x = F.relu(self.ra2_conv2(x))\n        x = F.relu(self.ra2_conv3(x))\n        ra2_feat = self.ra2_conv4(x)\n        x = ra2_feat + crop_2\n        lateral_map_2 = F.interpolate(x, scale_factor=8, mode='bilinear')   # NOTES: Sup-4 (bs, 1, 44, 44) -> (bs, 1, 352, 352)\n\n        return lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport numpy as np\n\ndef clip_gradient(optimizer, grad_clip):\n    \"\"\"\n    For calibrating misalignment gradient via cliping gradient technique\n    :param optimizer:\n    :param grad_clip:\n    :return:\n    \"\"\"\n    for group in optimizer.param_groups:\n        for param in group['params']:\n            if param.grad is not None:\n                param.grad.data.clamp_(-grad_clip, grad_clip)\n\n\ndef adjust_lr(optimizer, init_lr, epoch, decay_rate=0.1, decay_epoch=30):\n    decay = decay_rate ** (epoch // decay_epoch)\n    for param_group in optimizer.param_groups:\n        param_group['lr'] *= decay\n\n\nclass AvgMeter(object):\n    def __init__(self, num=40):\n        self.num = num\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n        self.losses = []\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n        self.losses.append(val)\n\n    def show(self):\n        return torch.mean(torch.stack(self.losses[np.maximum(len(self.losses)-self.num, 0):]))\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DiceLoss(nn.Module):\n    def __init__(self, weight=None, size_average=True):\n        super(DiceLoss, self).__init__()\n\n    def forward(self, inputs, targets, smooth=1):\n        \n        #comment out if your model contains a sigmoid or equivalent activation layer\n        inputs = torch.sigmoid(inputs)       \n        \n        #flatten label and prediction tensors\n        inputs = inputs.view(-1)\n        targets = targets.view(-1)\n        \n        intersection = (inputs * targets).sum()                            \n        dice = (2.*intersection + smooth)/(inputs.sum() + targets.sum() + smooth)  \n        \n        return 1 - dice\n    \n    \nloss_fn = DiceLoss()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nfrom torch.autograd import Variable\nimport os\nimport argparse\nfrom datetime import datetime\nimport torch.nn.functional as F\n\n\ndef structure_loss(pred, mask):\n    weit = 1 + 5*torch.abs(F.avg_pool2d(mask, kernel_size=31, stride=1, padding=15) - mask)\n    wbce = F.binary_cross_entropy_with_logits(pred, mask, reduce='none')\n    wbce = (weit*wbce).sum(dim=(2, 3)) / weit.sum(dim=(2, 3))\n\n    pred = torch.sigmoid(pred)\n    inter = ((pred * mask)*weit).sum(dim=(2, 3))\n    union = ((pred + mask)*weit).sum(dim=(2, 3))\n    wiou = 1 - (inter + 1)/(union - inter+1)\n    return (wbce + wiou).mean()\n\n\ndef train(train_loader, model, optimizer, epoch):\n    \n    model.train()\n    # ---- multi-scale training ----\n    size_rates = [0.75, 1, 1.25]\n    loss_record2, loss_record3, loss_record4, loss_record5, dice_loss_ = AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter()\n    for i, pack in enumerate(train_loader, start=1):\n        for rate in size_rates:\n            optimizer.zero_grad()\n            # ---- data prepare ----\n            images, gts = pack\n            images = Variable(images).type(torch.FloatTensor).to(device)\n            gts = Variable(gts).type(torch.FloatTensor).to(device)\n#             print(images.size(), gts.size())\n            # ---- rescale ----\n            trainsize = int(round(512*rate/32)*32)\n#             print(trainsize)\n            if rate != 1:\n                images = F.upsample(images, size=(trainsize, trainsize), mode='bilinear', align_corners=True)\n#                 print(images.size())\n                gts = F.upsample(gts, size=(trainsize, trainsize), mode='bilinear', align_corners=True)\n            # ---- forward ----\n            lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2 = model(images)\n            # ---- loss function ----\n            loss5 = structure_loss(lateral_map_5, gts)\n            loss4 = structure_loss(lateral_map_4, gts)\n            loss3 = structure_loss(lateral_map_3, gts)\n            loss2 = structure_loss(lateral_map_2, gts)\n            dl = loss_fn(lateral_map_2, gts)\n            loss = loss2 + loss3 + loss4 + loss5    # TODO: try different weights for loss\n            # ---- backward ----\n            loss.backward()\n            clip_gradient(optimizer, 0.5)\n            optimizer.step()\n            # ---- recording loss ----\n            if rate == 1:\n                dice_loss_.update(dl.data,  batchsize)\n                loss_record2.update(loss2.data, batchsize)\n                loss_record3.update(loss3.data, batchsize)\n                loss_record4.update(loss4.data, batchsize)\n                loss_record5.update(loss5.data, batchsize)\n        # ---- train visualization ----\n        if i % 30 == 0 or i == total_step:\n            print('{} Epoch [{:03d}/{:03d}], Step [{:04d}/{:04d}], '\n                  '[lateral-2: {:.4f}, lateral-3: {:0.4f}, lateral-4: {:0.4f}, lateral-5: {:0.4f}, Dice: {:0.4f}]'.\n                  format(datetime.now(), epoch, 10, i, total_step,\n                         loss_record2.show(), loss_record3.show(), loss_record4.show(), loss_record5.show(), dice_loss_.show()))\n    torch.save(model.state_dict(), 'PraNet-latest.pth' )\n    print('[Saving Snapshot:]', 'PraNet-latest.pth')\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nepoch = 10\ndecay_rate = 0.1\nlr = 0.001\ndecay_epoch = 2\ndevice = 'cuda'\n\nprint('Building Model...')\nmodel = PraNet()\nmodel.to(device)\n\nprint('Building Optimizer...')\noptimizer = torch.optim.Adam(model.parameters(), lr)\n\ntotal_step = len(train_loader)\n\nprint(\"#\"*20, \"Start Training\", \"#\"*20)\n\nfor epoch in range(1, epoch):\n    adjust_lr(optimizer, lr, epoch, decay_rate, decay_epoch)\n    train(train_loader, model, optimizer, epoch)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}