{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import copy\nimport os\nimport gc\nimport time\nimport random\nfrom abc import ABC\nfrom collections import defaultdict\nimport wandb\nimport skimage.morphology\n\nimport cv2 as cv\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport torch.utils.data\nimport torch.nn as nn\nfrom torch.cuda import amp\nfrom torch.optim import lr_scheduler\nfrom tqdm import tqdm\n\n\nfrom torch.utils.data import DataLoader\nimport sklearn.model_selection\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# EfficientNetEncoder-U-NET模型","metadata":{}},{"cell_type":"code","source":"import re\nimport math\nimport collections\nfrom functools import partial\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torch.utils import model_zoo\nfrom typing import Optional, Union, List\n\n########################################################################\n############### HELPERS FUNCTIONS FOR MODEL ARCHITECTURE ###############\n########################################################################\n\n\n# Parameters for the entire model (stem, all blocks, and head)\nGlobalParams = collections.namedtuple('GlobalParams', [\n    'batch_norm_momentum', 'batch_norm_epsilon', 'dropout_rate',\n    'num_classes', 'width_coefficient', 'depth_coefficient',\n    'depth_divisor', 'min_depth', 'drop_connect_rate', 'image_size'])\n\n# Parameters for an individual model block\nBlockArgs = collections.namedtuple('BlockArgs', [\n    'kernel_size', 'num_repeat', 'input_filters', 'output_filters',\n    'expand_ratio', 'id_skip', 'stride', 'se_ratio'])\n\n# Change namedtuple defaults\nGlobalParams.__new__.__defaults__ = (None,) * len(GlobalParams._fields)\nBlockArgs.__new__.__defaults__ = (None,) * len(BlockArgs._fields)\n\n\nclass SwishImplementation(torch.autograd.Function):\n    @staticmethod\n    def forward(ctx, i):\n        result = i * torch.sigmoid(i)\n        ctx.save_for_backward(i)\n        return result\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        i = ctx.saved_variables[0]\n        sigmoid_i = torch.sigmoid(i)\n        return grad_output * (sigmoid_i * (1 + i * (1 - sigmoid_i)))\n\n\nclass MemoryEfficientSwish(nn.Module):\n    def forward(self, x):\n        return SwishImplementation.apply(x)\n\nclass Swish(nn.Module):\n    def forward(self, x):\n        return x * torch.sigmoid(x)\n\n\ndef round_filters(filters, global_params):\n    \"\"\" Calculate and round number of filters based on depth multiplier. \"\"\"\n    multiplier = global_params.width_coefficient\n    if not multiplier:\n        return filters\n    divisor = global_params.depth_divisor\n    min_depth = global_params.min_depth\n    filters *= multiplier\n    min_depth = min_depth or divisor\n    new_filters = max(min_depth, int(filters + divisor / 2) // divisor * divisor)\n    if new_filters < 0.9 * filters:  # prevent rounding by more than 10%\n        new_filters += divisor\n    return int(new_filters)\n\n\ndef round_repeats(repeats, global_params):\n    \"\"\" Round number of filters based on depth multiplier. \"\"\"\n    multiplier = global_params.depth_coefficient\n    if not multiplier:\n        return repeats\n    return int(math.ceil(multiplier * repeats))\n\n\ndef drop_connect(inputs, p, training):\n    \"\"\" Drop connect. \"\"\"\n    if not training: return inputs\n    batch_size = inputs.shape[0]\n    keep_prob = 1 - p\n    random_tensor = keep_prob\n    random_tensor += torch.rand([batch_size, 1, 1, 1], dtype=inputs.dtype, device=inputs.device)\n    binary_tensor = torch.floor(random_tensor)\n    output = inputs / keep_prob * binary_tensor\n    return output\n\n\ndef get_same_padding_conv2d(image_size=None):\n    \"\"\" Chooses static padding if you have specified an image size, and dynamic padding otherwise.\n        Static padding is necessary for ONNX exporting of models. \"\"\"\n    if image_size is None:\n        return Conv2dDynamicSamePadding\n    else:\n        return partial(Conv2dStaticSamePadding, image_size=image_size)\n\n\nclass Conv2dDynamicSamePadding(nn.Conv2d):\n    \"\"\" 2D Convolutions like TensorFlow, for a dynamic image size \"\"\"\n\n    def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, groups=1, bias=True):\n        super().__init__(in_channels, out_channels, kernel_size, stride, 0, dilation, groups, bias)\n        self.stride = self.stride if len(self.stride) == 2 else [self.stride[0]] * 2\n\n    def forward(self, x):\n        ih, iw = x.size()[-2:]\n        kh, kw = self.weight.size()[-2:]\n        sh, sw = self.stride\n        oh, ow = math.ceil(ih / sh), math.ceil(iw / sw)\n        pad_h = max((oh - 1) * self.stride[0] + (kh - 1) * self.dilation[0] + 1 - ih, 0)\n        pad_w = max((ow - 1) * self.stride[1] + (kw - 1) * self.dilation[1] + 1 - iw, 0)\n        if pad_h > 0 or pad_w > 0:\n            x = F.pad(x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2])\n        return F.conv2d(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)\n\n\nclass Conv2dStaticSamePadding(nn.Conv2d):\n    \"\"\" 2D Convolutions like TensorFlow, for a fixed image size\"\"\"\n\n    def __init__(self, in_channels, out_channels, kernel_size, image_size=None, **kwargs):\n        super().__init__(in_channels, out_channels, kernel_size, **kwargs)\n        self.stride = self.stride if len(self.stride) == 2 else [self.stride[0]] * 2\n\n        # Calculate padding based on image size and save it\n        assert image_size is not None\n        ih, iw = image_size if type(image_size) == list else [image_size, image_size]\n        kh, kw = self.weight.size()[-2:]\n        sh, sw = self.stride\n        oh, ow = math.ceil(ih / sh), math.ceil(iw / sw)\n        pad_h = max((oh - 1) * self.stride[0] + (kh - 1) * self.dilation[0] + 1 - ih, 0)\n        pad_w = max((ow - 1) * self.stride[1] + (kw - 1) * self.dilation[1] + 1 - iw, 0)\n        if pad_h > 0 or pad_w > 0:\n            self.static_padding = nn.ZeroPad2d((pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2))\n        else:\n            self.static_padding = Identity()\n\n    def forward(self, x):\n        x = self.static_padding(x)\n        x = F.conv2d(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)\n        return x\n\n\nclass Identity(nn.Module):\n    def __init__(self, ):\n        super(Identity, self).__init__()\n\n    def forward(self, input):\n        return input\n\n\n########################################################################\n############## HELPERS FUNCTIONS FOR LOADING MODEL PARAMS ##############\n########################################################################\n\n\ndef efficientnet_params(model_name):\n    \"\"\" Map EfficientNet model name to parameter coefficients. \"\"\"\n    params_dict = {\n        # Coefficients:   width,depth,res,dropout\n        'efficientnet-b0': (1.0, 1.0, 224, 0.2),\n        'efficientnet-b1': (1.0, 1.1, 240, 0.2),\n        'efficientnet-b2': (1.1, 1.2, 260, 0.3),\n        'efficientnet-b3': (1.2, 1.4, 300, 0.3),\n        'efficientnet-b4': (1.4, 1.8, 380, 0.4),\n        'efficientnet-b5': (1.6, 2.2, 456, 0.4),\n        'efficientnet-b6': (1.8, 2.6, 528, 0.5),\n        'efficientnet-b7': (2.0, 3.1, 600, 0.5),\n        'efficientnet-b8': (2.2, 3.6, 672, 0.5),\n        'efficientnet-l2': (4.3, 5.3, 800, 0.5),\n    }\n    return params_dict[model_name]\n\n\nclass BlockDecoder(object):\n    \"\"\" Block Decoder for readability, straight from the official TensorFlow repository \"\"\"\n\n    @staticmethod\n    def _decode_block_string(block_string):\n        \"\"\" Gets a block through a string notation of arguments. \"\"\"\n        assert isinstance(block_string, str)\n\n        ops = block_string.split('_')\n        options = {}\n        for op in ops:\n            splits = re.split(r'(\\d.*)', op)\n            if len(splits) >= 2:\n                key, value = splits[:2]\n                options[key] = value\n\n        # Check stride\n        assert (('s' in options and len(options['s']) == 1) or\n                (len(options['s']) == 2 and options['s'][0] == options['s'][1]))\n\n        return BlockArgs(\n            kernel_size=int(options['k']),\n            num_repeat=int(options['r']),\n            input_filters=int(options['i']),\n            output_filters=int(options['o']),\n            expand_ratio=int(options['e']),\n            id_skip=('noskip' not in block_string),\n            se_ratio=float(options['se']) if 'se' in options else None,\n            stride=[int(options['s'][0])])\n\n    @staticmethod\n    def _encode_block_string(block):\n        \"\"\"Encodes a block to a string.\"\"\"\n        args = [\n            'r%d' % block.num_repeat,\n            'k%d' % block.kernel_size,\n            's%d%d' % (block.strides[0], block.strides[1]),\n            'e%s' % block.expand_ratio,\n            'i%d' % block.input_filters,\n            'o%d' % block.output_filters\n        ]\n        if 0 < block.se_ratio <= 1:\n            args.append('se%s' % block.se_ratio)\n        if block.id_skip is False:\n            args.append('noskip')\n        return '_'.join(args)\n\n    @staticmethod\n    def decode(string_list):\n        \"\"\"\n        Decodes a list of string notations to specify blocks inside the network.\n\n        :param string_list: a list of strings, each string is a notation of block\n        :return: a list of BlockArgs namedtuples of block args\n        \"\"\"\n        assert isinstance(string_list, list)\n        blocks_args = []\n        for block_string in string_list:\n            blocks_args.append(BlockDecoder._decode_block_string(block_string))\n        return blocks_args\n\n    @staticmethod\n    def encode(blocks_args):\n        \"\"\"\n        Encodes a list of BlockArgs to a list of strings.\n\n        :param blocks_args: a list of BlockArgs namedtuples of block args\n        :return: a list of strings, each string is a notation of block\n        \"\"\"\n        block_strings = []\n        for block in blocks_args:\n            block_strings.append(BlockDecoder._encode_block_string(block))\n        return block_strings\n\n\ndef efficientnet(width_coefficient=None, depth_coefficient=None, dropout_rate=0.2,\n                 drop_connect_rate=0.2, image_size=None, num_classes=1000):\n    \"\"\" Creates a efficientnet model. \"\"\"\n\n    blocks_args = [\n        'r1_k3_s11_e1_i32_o16_se0.25', 'r2_k3_s22_e6_i16_o24_se0.25',\n        'r2_k5_s22_e6_i24_o40_se0.25', 'r3_k3_s22_e6_i40_o80_se0.25',\n        'r3_k5_s11_e6_i80_o112_se0.25', 'r4_k5_s22_e6_i112_o192_se0.25',\n        'r1_k3_s11_e6_i192_o320_se0.25',\n    ]\n    blocks_args = BlockDecoder.decode(blocks_args)\n\n    global_params = GlobalParams(\n        batch_norm_momentum=0.99,\n        batch_norm_epsilon=1e-3,\n        dropout_rate=dropout_rate,\n        drop_connect_rate=drop_connect_rate,\n        # data_format='channels_last',  # removed, this is always true in PyTorch\n        num_classes=num_classes,\n        width_coefficient=width_coefficient,\n        depth_coefficient=depth_coefficient,\n        depth_divisor=8,\n        min_depth=None,\n        image_size=image_size,\n    )\n\n    return blocks_args, global_params\n\n\ndef get_model_params(model_name, override_params):\n    \"\"\" Get the block args and global params for a given model \"\"\"\n    if model_name.startswith('efficientnet'):\n        w, d, s, p = efficientnet_params(model_name)\n        # note: all models have drop connect rate = 0.2\n        blocks_args, global_params = efficientnet(\n            width_coefficient=w, depth_coefficient=d, dropout_rate=p, image_size=s)\n    else:\n        raise NotImplementedError('model name is not pre-defined: %s' % model_name)\n    if override_params:\n        # ValueError will be raised here if override_params has fields not included in global_params.\n        global_params = global_params._replace(**override_params)\n    return blocks_args, global_params\n\n\nurl_map = {\n    'efficientnet-b0': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b0-355c32eb.pth',\n    'efficientnet-b1': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b1-f1951068.pth',\n    'efficientnet-b2': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b2-8bb594d6.pth',\n    'efficientnet-b3': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b3-5fb5a3c3.pth',\n    'efficientnet-b4': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b4-6ed6700e.pth',\n    'efficientnet-b5': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b5-b6417697.pth',\n    'efficientnet-b6': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b6-c76e70fd.pth',\n    'efficientnet-b7': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b7-dcc49843.pth',\n}\n\n\nurl_map_advprop = {\n    'efficientnet-b0': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b0-b64d5a18.pth',\n    'efficientnet-b1': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b1-0f3ce85a.pth',\n    'efficientnet-b2': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b2-6e9d97e5.pth',\n    'efficientnet-b3': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b3-cdd7c0f4.pth',\n    'efficientnet-b4': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b4-44fb3a87.pth',\n    'efficientnet-b5': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b5-86493f6b.pth',\n    'efficientnet-b6': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b6-ac80338e.pth',\n    'efficientnet-b7': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b7-4652b6dd.pth',\n    'efficientnet-b8': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b8-22a8fe65.pth',\n}\n\n\ndef load_pretrained_weights(model, model_name, load_fc=True, advprop=False):\n    \"\"\" Loads pretrained weights, and downloads if loading for the first time. \"\"\"\n    # AutoAugment or Advprop (different preprocessing)\n    url_map_ = url_map_advprop if advprop else url_map\n    state_dict = model_zoo.load_url(url_map_[model_name])\n    if load_fc:\n        model.load_state_dict(state_dict)\n    else:\n        state_dict.pop('_fc.weight')\n        state_dict.pop('_fc.bias')\n        res = model.load_state_dict(state_dict, strict=False)\n        assert set(res.missing_keys) == set(['_fc.weight', '_fc.bias']), 'issue loading pretrained weights'\n    print('Loaded pretrained weights for {}'.format(model_name))\n\n\nclass MBConvBlock(nn.Module):\n    \"\"\"\n    Mobile Inverted Residual Bottleneck Block\n\n    Args:\n        block_args (namedtuple): BlockArgs, see above\n        global_params (namedtuple): GlobalParam, see above\n\n    Attributes:\n        has_se (bool): Whether the block contains a Squeeze and Excitation layer.\n    \"\"\"\n\n    def __init__(self, block_args, global_params):\n        super().__init__()\n        self._block_args = block_args\n        self._bn_mom = 1 - global_params.batch_norm_momentum\n        self._bn_eps = global_params.batch_norm_epsilon\n        self.has_se = (self._block_args.se_ratio is not None) and (0 < self._block_args.se_ratio <= 1)\n        self.id_skip = block_args.id_skip  # skip connection and drop connect\n\n        # Get static or dynamic convolution depending on image size\n        Conv2d = get_same_padding_conv2d(image_size=global_params.image_size)\n\n        # Expansion phase\n        inp = self._block_args.input_filters  # number of input channels\n        oup = self._block_args.input_filters * self._block_args.expand_ratio  # number of output channels\n        if self._block_args.expand_ratio != 1:\n            self._expand_conv = Conv2d(in_channels=inp, out_channels=oup, kernel_size=1, bias=False)\n            self._bn0 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps)\n\n        # Depthwise convolution phase\n        k = self._block_args.kernel_size\n        s = self._block_args.stride\n        self._depthwise_conv = Conv2d(\n            in_channels=oup, out_channels=oup, groups=oup,  # groups makes it depthwise\n            kernel_size=k, stride=s, bias=False)\n        self._bn1 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps)\n\n        # Squeeze and Excitation layer, if desired\n        if self.has_se:\n            num_squeezed_channels = max(1, int(self._block_args.input_filters * self._block_args.se_ratio))\n            self._se_reduce = Conv2d(in_channels=oup, out_channels=num_squeezed_channels, kernel_size=1)\n            self._se_expand = Conv2d(in_channels=num_squeezed_channels, out_channels=oup, kernel_size=1)\n\n        # Output phase\n        final_oup = self._block_args.output_filters\n        self._project_conv = Conv2d(in_channels=oup, out_channels=final_oup, kernel_size=1, bias=False)\n        self._bn2 = nn.BatchNorm2d(num_features=final_oup, momentum=self._bn_mom, eps=self._bn_eps)\n        self._swish = MemoryEfficientSwish()\n\n    def forward(self, inputs, drop_connect_rate=None):\n        \"\"\"\n        :param inputs: input tensor\n        :param drop_connect_rate: drop connect rate (float, between 0 and 1)\n        :return: output of block\n        \"\"\"\n\n        # Expansion and Depthwise Convolution\n        x = inputs\n        if self._block_args.expand_ratio != 1:\n            x = self._swish(self._bn0(self._expand_conv(inputs)))\n        x = self._swish(self._bn1(self._depthwise_conv(x)))\n\n        # Squeeze and Excitation\n        if self.has_se:\n            x_squeezed = F.adaptive_avg_pool2d(x, 1)\n            x_squeezed = self._se_expand(self._swish(self._se_reduce(x_squeezed)))\n            x = torch.sigmoid(x_squeezed) * x\n\n        x = self._bn2(self._project_conv(x))\n\n        # Skip connection and drop connect\n        input_filters, output_filters = self._block_args.input_filters, self._block_args.output_filters\n        if self.id_skip and self._block_args.stride == 1 and input_filters == output_filters:\n            if drop_connect_rate:\n                x = drop_connect(x, p=drop_connect_rate, training=self.training)\n            x = x + inputs  # skip connection\n        return x\n\n    def set_swish(self, memory_efficient=True):\n        \"\"\"Sets swish function as memory efficient (for training) or standard (for export)\"\"\"\n        self._swish = MemoryEfficientSwish() if memory_efficient else Swish()\n\n\nclass EfficientNet(nn.Module):\n    \"\"\"\n    An EfficientNet model. Most easily loaded with the .from_name or .from_pretrained methods\n\n    Args:\n        blocks_args (list): A list of BlockArgs to construct blocks\n        global_params (namedtuple): A set of GlobalParams shared between blocks\n\n    Example:\n        model = EfficientNet.from_pretrained('efficientnet-b0')\n\n    \"\"\"\n\n    def __init__(self, blocks_args=None, global_params=None):\n        super().__init__()\n        assert isinstance(blocks_args, list), 'blocks_args should be a list'\n        assert len(blocks_args) > 0, 'block args must be greater than 0'\n        self._global_params = global_params\n        self._blocks_args = blocks_args\n\n        # Get static or dynamic convolution depending on image size\n        Conv2d = get_same_padding_conv2d(image_size=global_params.image_size)\n\n        # Batch norm parameters\n        bn_mom = 1 - self._global_params.batch_norm_momentum\n        bn_eps = self._global_params.batch_norm_epsilon\n\n        # Stem\n        in_channels = 3  # rgb\n        out_channels = round_filters(32, self._global_params)  # number of output channels\n        self._conv_stem = Conv2d(in_channels, out_channels, kernel_size=3, stride=2, bias=False)\n        self._bn0 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps)\n\n        # Build blocks\n        self._blocks = nn.ModuleList([])\n        for block_args in self._blocks_args:\n\n            # Update block input and output filters based on depth multiplier.\n            block_args = block_args._replace(\n                input_filters=round_filters(block_args.input_filters, self._global_params),\n                output_filters=round_filters(block_args.output_filters, self._global_params),\n                num_repeat=round_repeats(block_args.num_repeat, self._global_params)\n            )\n\n            # The first block needs to take care of stride and filter size increase.\n            self._blocks.append(MBConvBlock(block_args, self._global_params))\n            if block_args.num_repeat > 1:\n                block_args = block_args._replace(input_filters=block_args.output_filters, stride=1)\n            for _ in range(block_args.num_repeat - 1):\n                self._blocks.append(MBConvBlock(block_args, self._global_params))\n\n        # Head\n        in_channels = block_args.output_filters  # output of final block\n        out_channels = round_filters(1280, self._global_params)\n        self._conv_head = Conv2d(in_channels, out_channels, kernel_size=1, bias=False)\n        self._bn1 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps)\n\n        # Final linear layer\n        self._avg_pooling = nn.AdaptiveAvgPool2d(1)\n        self._dropout = nn.Dropout(self._global_params.dropout_rate)\n        self._fc = nn.Linear(out_channels, self._global_params.num_classes)\n        self._swish = MemoryEfficientSwish()\n\n    def set_swish(self, memory_efficient=True):\n        \"\"\"Sets swish function as memory efficient (for training) or standard (for export)\"\"\"\n        self._swish = MemoryEfficientSwish() if memory_efficient else Swish()\n        for block in self._blocks:\n            block.set_swish(memory_efficient)\n\n    def extract_features(self, inputs):\n        \"\"\" Returns output of the final convolution layer \"\"\"\n\n        # Stem\n        x = self._swish(self._bn0(self._conv_stem(inputs)))\n\n        # Blocks\n        for idx, block in enumerate(self._blocks):\n            drop_connect_rate = self._global_params.drop_connect_rate\n            if drop_connect_rate:\n                drop_connect_rate *= float(idx) / len(self._blocks)\n            x = block(x, drop_connect_rate=drop_connect_rate)\n\n        # Head\n        x = self._swish(self._bn1(self._conv_head(x)))\n\n        return x\n\n    def forward(self, inputs):\n        \"\"\" Calls extract_features to extract features, applies final linear layer, and returns logits. \"\"\"\n        bs = inputs.size(0)\n        # Convolution layers\n        x = self.extract_features(inputs)\n\n        # Pooling and final linear layer\n        x = self._avg_pooling(x)\n        x = x.view(bs, -1)\n        x = self._dropout(x)\n        x = self._fc(x)\n        return x\n\n    @classmethod\n    def from_name(cls, model_name, override_params=None):\n        cls._check_model_name_is_valid(model_name)\n        blocks_args, global_params = get_model_params(model_name, override_params)\n        return cls(blocks_args, global_params)\n\n    @classmethod\n    def from_pretrained(cls, model_name, advprop=False, num_classes=1000, in_channels=3):\n        model = cls.from_name(model_name, override_params={'num_classes': num_classes})\n        load_pretrained_weights(model, model_name, load_fc=(num_classes == 1000), advprop=advprop)\n        if in_channels != 3:\n            Conv2d = get_same_padding_conv2d(image_size=model._global_params.image_size)\n            out_channels = round_filters(32, model._global_params)\n            model._conv_stem = Conv2d(in_channels, out_channels, kernel_size=3, stride=2, bias=False)\n        return model\n\n    @classmethod\n    def get_image_size(cls, model_name):\n        cls._check_model_name_is_valid(model_name)\n        _, _, res, _ = efficientnet_params(model_name)\n        return res\n\n    @classmethod\n    def _check_model_name_is_valid(cls, model_name):\n        \"\"\" Validates model name. \"\"\"\n        valid_models = ['efficientnet-b' + str(i) for i in range(9)]\n        if model_name not in valid_models:\n            raise ValueError('model_name should be one of: ' + ', '.join(valid_models))\n\n\n# 工具函数\ndef patch_first_conv(model, new_in_channels, default_in_channels=3, pretrained=True):\n    \"\"\"Change first convolution layer input channels.\n    In case:\n        in_channels == 1 or in_channels == 2 -> reuse original weights\n        in_channels > 3 -> make random kaiming normal initialization\n    \"\"\"\n\n    # get first conv\n    for module in model.modules():\n        if isinstance(module, nn.Conv2d) and module.in_channels == default_in_channels:\n            break\n\n    weight = module.weight.detach()\n    module.in_channels = new_in_channels\n\n    if not pretrained:\n        module.weight = nn.parameter.Parameter(\n            torch.Tensor(module.out_channels, new_in_channels // module.groups, *module.kernel_size)\n        )\n        module.reset_parameters()\n\n    elif new_in_channels == 1:\n        new_weight = weight.sum(1, keepdim=True)\n        module.weight = nn.parameter.Parameter(new_weight)\n\n    else:\n        new_weight = torch.Tensor(module.out_channels, new_in_channels // module.groups, *module.kernel_size)\n\n        for i in range(new_in_channels):\n            new_weight[:, i] = weight[:, i % default_in_channels]\n\n        new_weight = new_weight * (default_in_channels / new_in_channels)\n        module.weight = nn.parameter.Parameter(new_weight)\n\n\n# 工具函数\ndef replace_strides_with_dilation(module, dilation_rate):\n    \"\"\"Patch Conv2d modules replacing strides with dilation\"\"\"\n    for mod in module.modules():\n        if isinstance(mod, nn.Conv2d):\n            mod.stride = (1, 1)\n            mod.dilation = (dilation_rate, dilation_rate)\n            kh, kw = mod.kernel_size\n            mod.padding = ((kh // 2) * dilation_rate, (kh // 2) * dilation_rate)\n\n            # Kostyl for EfficientNet\n            if hasattr(mod, \"static_padding\"):\n                mod.static_padding = nn.Identity()\n\n\n# 由编码器继承，一些工具函数\nclass EncoderMixin:\n    \"\"\"Add encoder functionality such as:\n    - output channels specification of feature tensors (produced by encoder)\n    - patching first convolution for arbitrary input channels\n    \"\"\"\n\n    _output_stride = 32\n\n    @property\n    def out_channels(self):\n        \"\"\"Return channels dimensions for each tensor of forward output of encoder\"\"\"\n        return self._out_channels[: self._depth + 1]\n\n    @property\n    def output_stride(self):\n        return min(self._output_stride, 2 ** self._depth)\n\n    def set_in_channels(self, in_channels, pretrained=True):\n        \"\"\"Change first convolution channels\"\"\"\n        if in_channels == 3:\n            return\n\n        self._in_channels = in_channels\n        if self._out_channels[0] == 3:\n            self._out_channels = tuple([in_channels] + list(self._out_channels)[1:])\n\n        patch_first_conv(model=self, new_in_channels=in_channels, pretrained=pretrained)\n\n    def get_stages(self):\n        \"\"\"Override it in your implementation\"\"\"\n        raise NotImplementedError\n\n    def make_dilated(self, output_stride):\n\n        if output_stride == 16:\n            stage_list = [\n                5,\n            ]\n            dilation_list = [\n                2,\n            ]\n\n        elif output_stride == 8:\n            stage_list = [4, 5]\n            dilation_list = [2, 4]\n\n        else:\n            raise ValueError(\"Output stride should be 16 or 8, got {}.\".format(output_stride))\n\n        self._output_stride = output_stride\n\n        stages = self.get_stages()\n        for stage_indx, dilation_rate in zip(stage_list, dilation_list):\n            replace_strides_with_dilation(\n                module=stages[stage_indx],\n                dilation_rate=dilation_rate,\n            )\n\n\n# EfficientNetEncoder骨干网络的类\nclass EfficientNetEncoder(EfficientNet, EncoderMixin):\n    def __init__(self, stage_idxs, out_channels, model_name, depth=5):\n\n        blocks_args, global_params = get_model_params(model_name, override_params=None)\n        super().__init__(blocks_args, global_params)\n\n        self._stage_idxs = stage_idxs\n        self._out_channels = out_channels\n        self._depth = depth\n        self._in_channels = 3\n\n        del self._fc\n\n    def get_stages(self):\n        return [\n            nn.Identity(),\n            nn.Sequential(self._conv_stem, self._bn0, self._swish),\n            self._blocks[: self._stage_idxs[0]],\n            self._blocks[self._stage_idxs[0] : self._stage_idxs[1]],\n            self._blocks[self._stage_idxs[1] : self._stage_idxs[2]],\n            self._blocks[self._stage_idxs[2] :],\n        ]\n\n    def forward(self, x):\n        stages = self.get_stages()\n\n        block_number = 0.0\n        drop_connect_rate = self._global_params.drop_connect_rate\n\n        features = []\n        for i in range(self._depth + 1):\n\n            # Identity and Sequential stages\n            if i < 2:\n                x = stages[i](x)\n\n            # Block stages need drop_connect rate\n            else:\n                for module in stages[i]:\n                    drop_connect = drop_connect_rate * block_number / len(self._blocks)\n                    block_number += 1.0\n                    x = module(x, drop_connect)\n\n            features.append(x)\n\n        return features\n\n    def load_state_dict(self, state_dict, **kwargs):\n        state_dict.pop(\"_fc.bias\", None)\n        state_dict.pop(\"_fc.weight\", None)\n        super().load_state_dict(state_dict, **kwargs)\n\n\n# 获取预训练的各种设置\ndef _get_pretrained_settings(encoder):\n    pretrained_settings = {\n        \"imagenet\": {\n            \"mean\": [0.485, 0.456, 0.406],\n            \"std\": [0.229, 0.224, 0.225],\n            \"url\": url_map[encoder],\n            \"input_space\": \"RGB\",\n            \"input_range\": [0, 1],\n        },\n        \"advprop\": {\n            \"mean\": [0.5, 0.5, 0.5],\n            \"std\": [0.5, 0.5, 0.5],\n            \"url\": url_map_advprop[encoder],\n            \"input_space\": \"RGB\",\n            \"input_range\": [0, 1],\n        },\n    }\n    return pretrained_settings\n\n\n# 不同结构EfficientNetEncoder骨干网络的属性字典\nefficient_net_encoders = {\n    \"efficientnet-b0\": {\n        \"encoder\": EfficientNetEncoder,\n        \"pretrained_settings\": _get_pretrained_settings(\"efficientnet-b0\"),\n        \"params\": {\n            \"out_channels\": (3, 32, 24, 40, 112, 320),\n            \"stage_idxs\": (3, 5, 9, 16),\n            \"model_name\": \"efficientnet-b0\",\n        },\n    },\n    \"efficientnet-b1\": {\n        \"encoder\": EfficientNetEncoder,\n        \"pretrained_settings\": _get_pretrained_settings(\"efficientnet-b1\"),\n        \"params\": {\n            \"out_channels\": (3, 32, 24, 40, 112, 320),\n            \"stage_idxs\": (5, 8, 16, 23),\n            \"model_name\": \"efficientnet-b1\",\n        },\n    },\n    \"efficientnet-b2\": {\n        \"encoder\": EfficientNetEncoder,\n        \"pretrained_settings\": _get_pretrained_settings(\"efficientnet-b2\"),\n        \"params\": {\n            \"out_channels\": (3, 32, 24, 48, 120, 352),\n            \"stage_idxs\": (5, 8, 16, 23),\n            \"model_name\": \"efficientnet-b2\",\n        },\n    },\n    \"efficientnet-b3\": {\n        \"encoder\": EfficientNetEncoder,\n        \"pretrained_settings\": _get_pretrained_settings(\"efficientnet-b3\"),\n        \"params\": {\n            \"out_channels\": (3, 40, 32, 48, 136, 384),\n            \"stage_idxs\": (5, 8, 18, 26),\n            \"model_name\": \"efficientnet-b3\",\n        },\n    },\n    \"efficientnet-b4\": {\n        \"encoder\": EfficientNetEncoder,\n        \"pretrained_settings\": _get_pretrained_settings(\"efficientnet-b4\"),\n        \"params\": {\n            \"out_channels\": (3, 48, 32, 56, 160, 448),\n            \"stage_idxs\": (6, 10, 22, 32),\n            \"model_name\": \"efficientnet-b4\",\n        },\n    },\n    \"efficientnet-b5\": {\n        \"encoder\": EfficientNetEncoder,\n        \"pretrained_settings\": _get_pretrained_settings(\"efficientnet-b5\"),\n        \"params\": {\n            \"out_channels\": (3, 48, 40, 64, 176, 512),\n            \"stage_idxs\": (8, 13, 27, 39),\n            \"model_name\": \"efficientnet-b5\",\n        },\n    },\n    \"efficientnet-b6\": {\n        \"encoder\": EfficientNetEncoder,\n        \"pretrained_settings\": _get_pretrained_settings(\"efficientnet-b6\"),\n        \"params\": {\n            \"out_channels\": (3, 56, 40, 72, 200, 576),\n            \"stage_idxs\": (9, 15, 31, 45),\n            \"model_name\": \"efficientnet-b6\",\n        },\n    },\n    \"efficientnet-b7\": {\n        \"encoder\": EfficientNetEncoder,\n        \"pretrained_settings\": _get_pretrained_settings(\"efficientnet-b7\"),\n        \"params\": {\n            \"out_channels\": (3, 64, 48, 80, 224, 640),\n            \"stage_idxs\": (11, 18, 38, 55),\n            \"model_name\": \"efficientnet-b7\",\n        },\n    },\n}\n\n\n# 获取编码器的方法\ndef get_encoder(name, in_channels=3, depth=5, weights=None, output_stride=32):\n\n    try:\n        Encoder = efficient_net_encoders[name][\"encoder\"]\n    except KeyError:\n        raise KeyError(\"Wrong encoder name\")\n\n    params = efficient_net_encoders[name][\"params\"]\n    params.update(depth=depth)\n    encoder = Encoder(**params)\n\n    if weights is not None:\n        try:\n            settings = efficient_net_encoders[name][\"pretrained_settings\"][weights]\n        except KeyError:\n            raise KeyError(\"Wrong pretrained weights\")\n        encoder.load_state_dict(model_zoo.load_url(settings[\"url\"]))\n\n    encoder.set_in_channels(in_channels, pretrained=weights is not None)\n    if output_stride != 32:\n        encoder.make_dilated(output_stride)\n\n    return encoder\n\n\ntry:\n    from inplace_abn import InPlaceABN\nexcept ImportError:\n    InPlaceABN = None\n\n\n# 定义卷积模块的类\nclass Conv2dReLU(nn.Sequential):\n    def __init__(\n        self,\n        in_channels,\n        out_channels,\n        kernel_size,\n        padding=0,\n        stride=1,\n        use_batchnorm=True,\n    ):\n\n        if use_batchnorm == \"inplace\" and InPlaceABN is None:\n            raise RuntimeError(\n                \"In order to use `use_batchnorm='inplace'` inplace_abn package must be installed. \"\n                + \"To install see: https://github.com/mapillary/inplace_abn\"\n            )\n\n        conv = nn.Conv2d(\n            in_channels,\n            out_channels,\n            kernel_size,\n            stride=stride,\n            padding=padding,\n            bias=not (use_batchnorm),\n        )\n        relu = nn.ReLU(inplace=True)\n\n        if use_batchnorm == \"inplace\":\n            bn = InPlaceABN(out_channels, activation=\"leaky_relu\", activation_param=0.0)\n            relu = nn.Identity()\n\n        elif use_batchnorm and use_batchnorm != \"inplace\":\n            bn = nn.BatchNorm2d(out_channels)\n\n        else:\n            bn = nn.Identity()\n\n        super(Conv2dReLU, self).__init__(conv, bn, relu)\n\n\n# 定义SCSE激活层的类\nclass SCSEModule(nn.Module):\n    def __init__(self, in_channels, reduction=16):\n        super().__init__()\n        self.cSE = nn.Sequential(\n            nn.AdaptiveAvgPool2d(1),\n            nn.Conv2d(in_channels, in_channels // reduction, 1),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(in_channels // reduction, in_channels, 1),\n            nn.Sigmoid(),\n        )\n        self.sSE = nn.Sequential(nn.Conv2d(in_channels, 1, 1), nn.Sigmoid())\n\n    def forward(self, x):\n        return x * self.cSE(x) + x * self.sSE(x)\n\n\n# 定义ArgMax激活层的类\nclass ArgMax(nn.Module):\n    def __init__(self, dim=None):\n        super().__init__()\n        self.dim = dim\n\n    def forward(self, x):\n        return torch.argmax(x, dim=self.dim)\n\n\n# 定义Clamp激活层的类\nclass Clamp(nn.Module):\n    def __init__(self, min=0, max=1):\n        super().__init__()\n        self.min, self.max = min, max\n\n    def forward(self, x):\n        return torch.clamp(x, self.min, self.max)\n\n\n# 获取激活层的类\nclass Activation(nn.Module):\n    def __init__(self, name, **params):\n\n        super().__init__()\n\n        if name is None or name == \"identity\":\n            self.activation = nn.Identity(**params)\n        elif name == \"sigmoid\":\n            self.activation = nn.Sigmoid()\n        elif name == \"softmax2d\":\n            self.activation = nn.Softmax(dim=1, **params)\n        elif name == \"softmax\":\n            self.activation = nn.Softmax(**params)\n        elif name == \"logsoftmax\":\n            self.activation = nn.LogSoftmax(**params)\n        elif name == \"tanh\":\n            self.activation = nn.Tanh()\n        elif name == \"argmax\":\n            self.activation = ArgMax(**params)\n        elif name == \"argmax2d\":\n            self.activation = ArgMax(dim=1, **params)\n        elif name == \"clamp\":\n            self.activation = Clamp(**params)\n        elif callable(name):\n            self.activation = name(**params)\n        else:\n            raise ValueError(\n                f\"Activation should be callable/sigmoid/softmax/logsoftmax/tanh/\"\n                f\"argmax/argmax2d/clamp/None; got {name}\"\n            )\n\n    def forward(self, x):\n        return self.activation(x)\n\n\n# 定义注意力层？\nclass Attention(nn.Module):\n    def __init__(self, name, **params):\n        super().__init__()\n\n        if name is None:\n            self.attention = nn.Identity(**params)\n        elif name == \"scse\":\n            self.attention = SCSEModule(**params)\n        else:\n            raise ValueError(\"Attention {} is not implemented\".format(name))\n\n    def forward(self, x):\n        return self.attention(x)\n\n\n# 定义分割头的类\nclass SegmentationHead(nn.Sequential):\n    def __init__(self, in_channels, out_channels, kernel_size=3, activation=None, upsampling=1):\n        conv2d = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2)\n        upsampling = nn.UpsamplingBilinear2d(scale_factor=upsampling) if upsampling > 1 else nn.Identity()\n        activation = Activation(activation)\n        super().__init__(conv2d, upsampling, activation)\n\n\n# 定义分类头的类\nclass ClassificationHead(nn.Sequential):\n    def __init__(self, in_channels, classes, pooling=\"avg\", dropout=0.2, activation=None):\n        if pooling not in (\"max\", \"avg\"):\n            raise ValueError(\"Pooling should be one of ('max', 'avg'), got {}.\".format(pooling))\n        pool = nn.AdaptiveAvgPool2d(1) if pooling == \"avg\" else nn.AdaptiveMaxPool2d(1)\n        flatten = nn.Flatten()\n        dropout = nn.Dropout(p=dropout, inplace=True) if dropout else nn.Identity()\n        linear = nn.Linear(in_channels, classes, bias=True)\n        activation = Activation(activation)\n        super().__init__(pool, flatten, dropout, linear, activation)\n\n\n# 初始化解码器的方法\ndef initialize_decoder(module):\n    for m in module.modules():\n\n        if isinstance(m, nn.Conv2d):\n            nn.init.kaiming_uniform_(m.weight, mode=\"fan_in\", nonlinearity=\"relu\")\n            if m.bias is not None:\n                nn.init.constant_(m.bias, 0)\n\n        elif isinstance(m, nn.BatchNorm2d):\n            nn.init.constant_(m.weight, 1)\n            nn.init.constant_(m.bias, 0)\n\n        elif isinstance(m, nn.Linear):\n            nn.init.xavier_uniform_(m.weight)\n            if m.bias is not None:\n                nn.init.constant_(m.bias, 0)\n\n\n# 初始化头模块的方法\ndef initialize_head(module):\n    for m in module.modules():\n        if isinstance(m, (nn.Linear, nn.Conv2d)):\n            nn.init.xavier_uniform_(m.weight)\n            if m.bias is not None:\n                nn.init.constant_(m.bias, 0)\n\n\n# 定义所有分割模型的父类\nclass SegmentationModel(torch.nn.Module):\n    def initialize(self):\n        initialize_decoder(self.decoder)\n        initialize_head(self.segmentation_head)\n        if self.classification_head is not None:\n            initialize_head(self.classification_head)\n\n    def check_input_shape(self, x):\n\n        h, w = x.shape[-2:]\n        output_stride = self.encoder.output_stride\n        if h % output_stride != 0 or w % output_stride != 0:\n            new_h = (h // output_stride + 1) * output_stride if h % output_stride != 0 else h\n            new_w = (w // output_stride + 1) * output_stride if w % output_stride != 0 else w\n            raise RuntimeError(\n                f\"Wrong input shape height={h}, width={w}. Expected image height and width \"\n                f\"divisible by {output_stride}. Consider pad your images to shape ({new_h}, {new_w}).\"\n            )\n\n    def forward(self, x):\n        \"\"\"Sequentially pass `x` trough model`s encoder, decoder and heads\"\"\"\n\n        self.check_input_shape(x)\n\n        features = self.encoder(x)\n        decoder_output = self.decoder(*features)\n\n        masks = self.segmentation_head(decoder_output)\n\n        if self.classification_head is not None:\n            labels = self.classification_head(features[-1])\n            return masks, labels\n\n        return masks\n\n    @torch.no_grad()\n    def predict(self, x):\n        \"\"\"Inference method. Switch model to `eval` mode, call `.forward(x)` with `torch.no_grad()`\n        Args:\n            x: 4D torch tensor with shape (batch_size, channels, height, width)\n        Return:\n            prediction: 4D torch tensor with shape (batch_size, classes, height, width)\n        \"\"\"\n        if self.training:\n            self.eval()\n\n        x = self.forward(x)\n\n        return x\n\n\n# 定义解码模块的类\nclass DecoderBlock(nn.Module):\n    def __init__(\n        self,\n        in_channels,\n        skip_channels,\n        out_channels,\n        use_batchnorm=True,\n        attention_type=None,\n    ):\n        super().__init__()\n        self.conv1 = Conv2dReLU(\n            in_channels + skip_channels,\n            out_channels,\n            kernel_size=3,\n            padding=1,\n            use_batchnorm=use_batchnorm,\n        )\n        self.attention1 = Attention(attention_type, in_channels=in_channels + skip_channels)\n        self.conv2 = Conv2dReLU(\n            out_channels,\n            out_channels,\n            kernel_size=3,\n            padding=1,\n            use_batchnorm=use_batchnorm,\n        )\n        self.attention2 = Attention(attention_type, in_channels=out_channels)\n\n    def forward(self, x, skip=None):\n        x = F.interpolate(x, scale_factor=2, mode=\"nearest\")\n        if skip is not None:\n            x = torch.cat([x, skip], dim=1)\n            x = self.attention1(x)\n        x = self.conv1(x)\n        x = self.conv2(x)\n        x = self.attention2(x)\n        return x\n\n\n# 定义中心模块的类\nclass CenterBlock(nn.Sequential):\n    def __init__(self, in_channels, out_channels, use_batchnorm=True):\n        conv1 = Conv2dReLU(\n            in_channels,\n            out_channels,\n            kernel_size=3,\n            padding=1,\n            use_batchnorm=use_batchnorm,\n        )\n        conv2 = Conv2dReLU(\n            out_channels,\n            out_channels,\n            kernel_size=3,\n            padding=1,\n            use_batchnorm=use_batchnorm,\n        )\n        super().__init__(conv1, conv2)\n\n\n# 定义整体U-Net解码器的类\nclass UnetDecoder(nn.Module):\n    def __init__(\n        self,\n        encoder_channels,\n        decoder_channels,\n        n_blocks=5,\n        use_batchnorm=True,\n        attention_type=None,\n        center=False,\n    ):\n        super().__init__()\n\n        if n_blocks != len(decoder_channels):\n            raise ValueError(\n                \"Model depth is {}, but you provide `decoder_channels` for {} blocks.\".format(\n                    n_blocks, len(decoder_channels)\n                )\n            )\n\n        # remove first skip with same spatial resolution\n        encoder_channels = encoder_channels[1:]\n        # reverse channels to start from head of encoder\n        encoder_channels = encoder_channels[::-1]\n\n        # computing blocks input and output channels\n        head_channels = encoder_channels[0]\n        in_channels = [head_channels] + list(decoder_channels[:-1])\n        skip_channels = list(encoder_channels[1:]) + [0]\n        out_channels = decoder_channels\n\n        if center:\n            self.center = CenterBlock(head_channels, head_channels, use_batchnorm=use_batchnorm)\n        else:\n            self.center = nn.Identity()\n\n        # combine decoder keyword arguments\n        kwargs = dict(use_batchnorm=use_batchnorm, attention_type=attention_type)\n        blocks = [\n            DecoderBlock(in_ch, skip_ch, out_ch, **kwargs)\n            for in_ch, skip_ch, out_ch in zip(in_channels, skip_channels, out_channels)\n        ]\n        self.blocks = nn.ModuleList(blocks)\n\n    def forward(self, *features):\n\n        features = features[1:]  # remove first skip with same spatial resolution\n        features = features[::-1]  # reverse channels to start from head of encoder\n\n        head = features[0]\n        skips = features[1:]\n\n        x = self.center(head)\n        for i, decoder_block in enumerate(self.blocks):\n            skip = skips[i] if i < len(skips) else None\n            x = decoder_block(x, skip)\n\n        return x\n\n\n# 结合EfficientNetEncoder骨干网络，定义整体U-Net网络\nclass EfficientNetEncoder_Unet(SegmentationModel):\n\n    def __init__(\n        self,\n        encoder_name: str = \"efficientnet-b0\",\n        encoder_depth: int = 5,\n        encoder_weights: Optional[str] = \"imagenet\",\n        decoder_use_batchnorm: bool = True,\n        decoder_channels: List[int] = (256, 128, 64, 32, 16),\n        decoder_attention_type: Optional[str] = None,\n        in_channels: int = 3,\n        classes: int = 1,\n        activation: Optional[Union[str, callable]] = None,\n        aux_params: Optional[dict] = None,\n    ):\n        super().__init__()\n\n        self.encoder = get_encoder(\n            encoder_name,\n            in_channels=in_channels,\n            depth=encoder_depth,\n            weights=encoder_weights,\n        )\n\n        self.decoder = UnetDecoder(\n            encoder_channels=self.encoder.out_channels,\n            decoder_channels=decoder_channels,\n            n_blocks=encoder_depth,\n            use_batchnorm=decoder_use_batchnorm,\n            center=True if encoder_name.startswith(\"vgg\") else False,\n            attention_type=decoder_attention_type,\n        )\n\n        self.segmentation_head = SegmentationHead(\n            in_channels=decoder_channels[-1],\n            out_channels=classes,\n            activation=activation,\n            kernel_size=3,\n        )\n\n        if aux_params is not None:\n            self.classification_head = ClassificationHead(in_channels=self.encoder.out_channels[-1], **aux_params)\n        else:\n            self.classification_head = None\n\n        self.name = \"u-{}\".format(encoder_name)\n        self.initialize()","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 损失函数","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn.functional as F\nfrom torch.nn.modules.loss import _Loss\n\n\nclass Dice(_Loss):\n    def __init__(self, smooth=0., eps=1e-7, dims=(0, 2), from_logic=True):\n        super().__init__()\n        self.smooth = smooth\n        self.eps = eps\n        self.dims = dims\n        self.from_logic = from_logic\n\n    def forward(self, y_pred, y_true):\n\n        if self.from_logic:\n            y_pred = F.logsigmoid(y_pred).exp()\n\n        bs = y_true.size(0)  # batch_size 每个tensor的第一维大小\n\n        # 将每个图像转化为一维（bs,1,512,512）--->（bs,1,512*512）\n        target = y_true.view(bs, 1, -1)\n        output = y_pred.view(bs, 1, -1)\n\n        # IOU score的计算\n        intersection = torch.sum(output * target, dim=self.dims)  # 计算交集面积\n        cardinality = torch.sum(output + target, dim=self.dims)  # 计算并集面积\n        score = (2*intersection + self.smooth) / (cardinality + self.smooth).clamp_min(self.eps)\n\n        # loss的计算\n        loss = 1 - score\n        mask = target.sum(self.dims) > 0\n        loss *= mask.float()\n\n        return loss.mean()\n\n    \nclass Jaccard(_Loss):\n    def __init__(self, smooth=0.0, eps=1e-7, dims=(0, 2), from_logic=True):\n        super(Jaccard, self).__init__()\n        self.smooth = smooth\n        self.eps = eps\n        self.dims = dims\n        self.from_logic = from_logic\n\n    def forward(self, y_pred, y_true) -> torch.Tensor:\n\n        if self.from_logic:\n            y_pred = F.logsigmoid(y_pred).exp()\n\n        bs = y_true.size(0)  # batch_size 每个tensor的第一维大小\n\n        # 将每个图像转化为一维（bs,1,512,512）--->（bs,1,512*512）\n        target = y_true.view(bs, 1, -1)\n        output = y_pred.view(bs, 1, -1)\n\n        # IOU score的计算\n        intersection = torch.sum(output * target, dim=self.dims)  # 计算交集面积\n        cardinality = torch.sum(output + target, dim=self.dims)  # 计算并集面积\n        union = cardinality - intersection  # 并集减去交集，作为IOU的分母\n        score = (intersection + self.smooth) / (union + self.smooth).clamp_min(self.eps)\n\n        # loss的计算\n        loss = 1.0 - score\n        mask = target.sum(self.dims) > 0\n        loss *= mask.float()\n\n        return loss.mean()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 定义超参数","metadata":{}},{"cell_type":"code","source":"class cfg:\n    seed          = 101\n    debug         = False # set debug=False for Full Training\n    exp_name      = 'Unet-effnetb2-512x512-aug2'\n    model_name    = 'Unet'\n    backbone      = 'vgg11'\n    train_bs      = 24\n    valid_bs      = 48\n    img_size      = [512, 512]\n    infer_size    = [520, 704]\n    epoch_multiple= 1\n    epochs        = 50*epoch_multiple\n    lr            = 5e-3\n    scheduler     = 'CosineAnnealingLR'\n    min_lr        = 1e-6\n    T_max         = int(100*6*1.8)*epoch_multiple\n    T_0           = 25*epoch_multiple\n    warmup_epochs = 0\n    wd            = 1e-6\n    n_accumulate  = 32//train_bs\n    n_fold        = 5\n    num_classes   = 1\n    device        = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n    competition   = 'sartorius'\n    _wandb_kernel = 'exiaoseiei'","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 定义图像增强pipeline","metadata":{}},{"cell_type":"code","source":"transforms_sartorius = {\n    'train': A.Compose([A.Resize(*cfg.img_size),\n                     # A.Normalize(mean=[0.485, 0.456, 0.406],\n                     #             std=[0.229, 0.224, 0.225],\n                     #             max_pixel_value=255.0,\n                     #             p=1.0),\n                     A.CLAHE(p=0.35),\n                     A.ColorJitter(p=0.5),\n                     A.HorizontalFlip(p=0.5),\n                     A.VerticalFlip(p=0.5),\n                     A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.1, rotate_limit=90, p=0.5),\n                     A.OneOf([A.GridDistortion(num_steps=5, distort_limit=0.05, p=1.0),\n                              A.ElasticTransform(alpha=1, sigma=50, alpha_affine=50, p=1.0)]),\n                     A.CoarseDropout(max_holes=8, max_height=cfg.img_size[0] // 20,\n                                     max_width=cfg.img_size[1] // 20,\n                                     min_holes=5, fill_value=0, mask_fill_value=0, p=0.5),\n                     ToTensorV2()\n                     ],\n                     p=1),\n\n    'valid': A.Compose([A.Resize(*cfg.img_size),\n                     # A.Normalize(mean=[0.485, 0.456, 0.406],\n                     #             std=[0.229, 0.224, 0.225],\n                     #             max_pixel_value=255.0,\n                     #             p=1.0),\n                     ToTensorV2()\n                     ],\n                     p=1),\n\n    'infer': A.Compose([A.Resize(*cfg.infer_size)],p=1)\n                        }","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 设置全局随机种子","metadata":{}},{"cell_type":"code","source":"def set_seed(seed):\n    np.random.seed(seed)\n    random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    # When running on the CuDNN backend, two further options must be set\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n    # Set a fixed value for the hash seed\n    os.environ['PYTHONHASHSEED'] = str(seed)\n\nset_seed(cfg.seed)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 一些工具函数","metadata":{}},{"cell_type":"code","source":"# 图片显示函数\ndef show_image(image):\n    cv.imshow('show', image)\n    cv.waitKey(0)\n\n# 将整张图片的RLE序列转为一个整体mask：\ndef rle2mask_list(rle_list, width=520, height=704):\n    mask_sum = np.zeros(width * height)\n    for rle in rle_list:\n        mask = np.zeros(width * height)\n        array = np.asarray([int(x) for x in rle.split(' ')])\n        starts = array[0::2]\n        lengths = array[1::2]\n        for index, start in enumerate(starts):\n            current_position = start\n            mask[current_position:current_position+lengths[index]] = 1\n            current_position += lengths[index]\n        mask_sum += mask\n    mask_sum[mask_sum > 0] = 1\n    return mask_sum.reshape(width, height)\n\n# 将单个实例的mask转化为RLE编码\ndef mask2rle(mask):\n    mask = np.array(mask)\n    pixels = mask.flatten()\n    pad = np.array([0])\n    pixels = np.concatenate([pad, pixels, pad])\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n    runs[1::2] -= runs[::2]\n    return ' '.join(str(x) for x in runs)\n\n# 将整张mask按连通区域分割为多个实例\ndef mask2rle_list(mask, cutoff=0.5, min_object_size=1.0):\n    \"\"\" Return run length encoding of mask.\n        ref: https://www.kaggle.com/raoulma/nuclei-dsb-2018-tensorflow-u-net-score-0-352\n    \"\"\"\n    # segment image and label different objects\n    lab_mask = skimage.morphology.label(mask > cutoff)\n\n    # Keep only objects that are large enough.\n    (mask_labels, mask_sizes) = np.unique(lab_mask, return_counts=True)\n    l = (mask_sizes < min_object_size).any()\n    if (mask_sizes < min_object_size).any():\n        mask_labels = mask_labels[mask_sizes < min_object_size]\n        for n in mask_labels:\n            lab_mask[lab_mask == n] = 0\n        lab_mask = skimage.morphology.label(lab_mask > cutoff)\n\n        # Loop over each object excluding the background labeled by 0.\n    for i in range(1, lab_mask.max() + 1):\n        yield mask2rle(lab_mask == i)\n\n# 初始化策略\ndef init_weights(m):\n    classname = m.__class__.__name__\n\n    # 卷积层初始化权重策略\n    if classname.find('Conv2d') != -1:\n        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')  # 凯明正态分布初始化\n\n    # 正则化层初始化权重策略\n    if classname.find('BatchNorm2d') != -1:\n        nn.init.constant_(m.weight, 1)  # 常数初始化\n        nn.init.constant_(m.bias, 0)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 数据准备函数","metadata":{}},{"cell_type":"code","source":"# 利用训练集数据，合成新dataframe，合并同一图片的所有mask编码\ndef get_df():\n    df = pd.read_csv('/kaggle/input/sartorius-cell-instance-segmentation/train.csv')\n    df['image_path'] = '/kaggle/input/sartorius-cell-instance-segmentation/train/' + df['id'] + '.png'\n    annotation_list = df.groupby('id')['annotation'].agg(list).reset_index(drop=True)\n    df_c = df.copy(deep=True)\n    df_c = df_c.drop_duplicates('id').reset_index(drop=True)\n    df_c['annotation'] = annotation_list\n    \n    # 构造K折交叉验证分选标签\n    skf = sklearn.model_selection.StratifiedKFold(n_splits=cfg.n_fold, shuffle=True, random_state=cfg.seed)\n    for K, (tra_index, val_index) in enumerate(skf.split(df_c, df_c['cell_type'])):\n        df_c.loc[val_index, 'fold'] = K\n    \n    return df_c\n\n# 自写dataset类\nclass Df2Dataset(torch.utils.data.Dataset):\n    def __init__(self, dataframe, transfroms=None):\n        self.dataframe = dataframe\n        self.img_path = dataframe['image_path']\n        try:\n            self.mask_rle = dataframe['annotation']\n        except:\n            self.mask_rle = None\n        self.transfroms = transfroms\n\n    def __len__(self):\n        return len(self.dataframe)\n\n    def __getitem__(self, item):\n        img = cv.imread(self.img_path[item])\n        img = cv.cvtColor(img, cv.COLOR_BGR2RGB)\n        if self.mask_rle is None:\n            if self.transfroms is not None:\n                tf_re = self.transfroms(image=img)\n                img = tf_re['image']\n            return img\n        if self.mask_rle is not None:\n            mask = rle2mask_list(self.mask_rle[item])\n            if self.transfroms is not None:\n                tf_re = self.transfroms(image=img, mask=mask)\n                img = tf_re['image']\n                mask = tf_re['mask']\n            mask = np.expand_dims(mask, axis=0)\n            return img, mask\n\n\n# 根据交叉实验返回相应的dataloader\ndef get_fold_dataloader(fold, dataframe):\n    train_df = dataframe[dataframe['fold'] != fold].reset_index(drop=True)\n    valid_df = dataframe[dataframe['fold'] == fold].reset_index(drop=True)\n\n    train_ds = Df2Dataset(train_df, transfroms=transforms_sartorius['train'])\n    valid_ds = Df2Dataset(valid_df, transfroms=transforms_sartorius['valid'])\n\n    train_dl = torch.utils.data.DataLoader(train_ds, batch_size=cfg.train_bs, shuffle=True, num_workers=0,\n                                           pin_memory=True)\n    valid_dl = torch.utils.data.DataLoader(valid_ds, batch_size=cfg.train_bs, shuffle=False, num_workers=0,\n                                           pin_memory=True)\n\n    return train_dl, valid_dl","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 定义模型，损失函数，学习率更新策略函数","metadata":{}},{"cell_type":"code","source":"# 创建模型函数，U-NET\ndef build_model():\n#     model = smp.Unet(encoder_name=cfg.backbone, encoder_weights='imagenet', in_channels=3, classes=cfg.num_classes,\n#                      activation=None)\n    model.to(cfg.device)  # 表示将模型加载到指定设备上\n    return model\n\n# 创建自写模型函数，U-NET\ndef build_model_my():\n    model = EfficientNetEncoder_Unet(in_channels=3, classes=cfg.num_classes,activation=None, encoder_weights=None, encoder_name='efficientnet-b3')\n#     model.apply(init_weights)  # 对模型使用指定权重初始化策略\n    model.to(cfg.device)  # 将模型加载到指定设备上\n    return model\n\n# 加载模型函数\ndef load_model(path):\n    model = build_model()\n    model.load_state_dict(torch.load(path))\n    model.eval()\n    return model\n\n# 定义损失函数\ndef criterion(y_pred, y_true):\n    JaccardLoss = Jaccard()\n    return JaccardLoss(y_pred, y_true)\n\n# 定义学习率更新策略，包装优化器\ndef get_scheduler(optimizer):\n    if cfg.scheduler == 'CosineAnnealingLR':\n        scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=cfg.T_max,\n                                                   eta_min=cfg.min_lr)\n    elif cfg.scheduler == 'CosineAnnealingWarmRestarts':\n        scheduler = lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=cfg.T_0,\n                                                             eta_min=cfg.min_lr)\n    elif cfg.scheduler == 'ReduceLROnPlateau':\n        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer,\n                                                   mode='min',\n                                                   factor=0.1,\n                                                   patience=7,\n                                                   threshold=0.0001,\n                                                   min_lr=cfg.min_lr, )\n    elif cfg.scheduler == 'ExponentialLR':\n        scheduler = lr_scheduler.ExponentialLR(optimizer, gamma=0.85)\n    elif cfg.scheduler == None:\n        return None\n\n    return scheduler","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 模型的one-epoch-runner","metadata":{}},{"cell_type":"code","source":"# 训练一个epoch的函数\ndef train_one_epoch(model, optimizer, scheduler, dataloader, epoch, device):\n    model.train()\n    # 单双精度转换策略\n    scaler = amp.GradScaler()\n\n    sum_loss = 0.0\n    sum_size = 0\n    mean_tra_loss = 0\n\n    # 训练一个epoch内的所有mini_batch\n    pbar = tqdm(enumerate(dataloader), total=len(dataloader), desc=f'正在训练第{epoch}轮')\n    for step, (images, masks) in pbar:\n        images = images.to(device, dtype=torch.float)\n        masks = masks.to(device, dtype=torch.float)\n        batch_size = images.shape[0]\n\n        # 开启单精度加速模式\n        with amp.autocast(enabled=True):\n            # 正向传播\n            y_pred = model(images)\n            # 计算损失，由于采用 n_accumulate 个 mini_batch 积累后求均值更新的策略，因此需要将损失除以 n_accumulate\n            loss = criterion(y_pred, masks)\n            loss = loss/cfg.n_accumulate\n\n        # 反向传播\n        scaler.scale(loss).backward()\n        # 是否达到mini_batch积累值\n        if (step + 1) % cfg.n_accumulate == 0:\n            scaler.step(optimizer)\n            scaler.update()\n            optimizer.zero_grad()\n\n            # 进行学习率调整\n            if scheduler is not None:\n                scheduler.step()\n\n        # 计算一个epoch的平均loss\n        sum_loss += loss.item()*batch_size\n        sum_size += batch_size\n        mean_tra_loss = sum_loss/sum_size\n        # 获取gpu占用率\n        mem = (torch.cuda.memory_reserved() / 1E9) if torch.cuda.is_available() else 0\n\n        # 设置进度条显示\n        pbar.set_postfix(train_loss=f'{mean_tra_loss:0.4f}',\n                         leaning_rate=optimizer.param_groups[0]['lr'],\n                         GPU_memory=f'{mem:0.2f}GB')\n\n    # 释放显存与内存\n    torch.cuda.empty_cache()\n    gc.collect()\n\n    return mean_tra_loss\n\n\n# 验证一个epoch的函数，整个过程不需要更新梯度\n@ torch.no_grad()\ndef valid_one_epoch(model, dataloader, epoch, device):\n    model.eval()\n\n    sum_loss = 0.0\n    sum_size = 0\n    mean_val_loss = 0\n\n    target_list = []\n    pred_list = []\n\n    pbar = tqdm(enumerate(dataloader), total=len(dataloader), desc=f'正在验证第{epoch}轮')\n    for step, (images, masks) in pbar:\n        images = images.to(device, dtype=torch.float)\n        masks = masks.to(device, dtype=torch.float)\n        batch_size = images.shape[0]\n\n        pred_y = model(images)\n        loss = criterion(pred_y, masks)\n\n        target_list.append(masks)\n        pred_list.append(nn.Sigmoid()(pred_y))\n\n        sum_loss += loss.item()*batch_size\n        sum_size += batch_size\n        mean_val_loss = sum_loss/sum_size\n        mem = (torch.cuda.memory_reserved()/1E9) if torch.cuda.is_available() else 0\n\n        pbar.set_postfix(valid_loss=f'{mean_val_loss:0.4f}',\n                         GPU_memory=f'{mem:0.2f}GB')\n\n    # 根据总输出mask与真实mask计算各项评价指标\n    targets = torch.cat(target_list, dim=0).to(dtype=torch.float32)\n    preds = (torch.cat(pred_list, dim=0) > 0.5).to(dtype=torch.float32)\n    \n    jaccard = Jaccard(from_logic=False)\n    dice = Dice(from_logic=False)\n\n    val_dice = 1 - dice(y_true=targets, y_pred=preds).cpu().detach().numpy()\n    val_jaccard = 1 - jaccard(y_true=targets, y_pred=preds).cpu().detach().numpy()\n    \n    val_scores = [val_dice, val_jaccard]\n\n    # 释放显存与内存\n    torch.cuda.empty_cache()\n    gc.collect()\n\n    return mean_val_loss, val_scores","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 模型的epoches-runner","metadata":{}},{"cell_type":"code","source":"# 运行多个epoch的函数\ndef run_epochs(model, train_dataloader, valid_dataloader, optimizer, scheduler, num_epoch, fold, save=False):\n\n    wandb.watch(model, log_freq=100)\n\n    if torch.cuda.is_available():\n        print(\"cuda: {}\\n\".format(torch.cuda.get_device_name()))\n\n    start_time = time.time()\n    best_dice = -np.inf  # 负无穷大\n    best_epoch = -1\n    best_model_wts = copy.deepcopy(model.state_dict())\n\n    history = defaultdict(list)\n\n    for epoch in range(1, num_epoch + 1):\n\n        gc.collect()\n\n        print(f'Epoch: {epoch}/{num_epoch}')\n        train_loss = train_one_epoch(model, optimizer, scheduler, train_dataloader, epoch, cfg.device)\n        valid_loss, val_scores = valid_one_epoch(model, valid_dataloader, epoch, cfg.device)\n\n        val_dice = val_scores[0]\n        val_jaccard = val_scores[1]\n\n        history['训练集损失'].append(train_loss)\n        history['验证集损失'].append(valid_loss)\n        history['验证集DICE'].append(val_dice)\n        history['验证集JACCARD'].append(val_jaccard)\n\n        # 上传WB数据\n        wandb.log({'训练集损失：': train_loss,\n                   '验证集损失：': valid_loss,\n                   '验证集DICE：': val_dice,\n                   '验证集JACCARD：': val_jaccard,\n                   '学习率：': scheduler.get_last_lr()[0]})\n\n        print(f'Val_Dice:{val_dice:0.4f} || Val_Jaccard:{val_jaccard:0.4f}')\n\n        if val_dice >= best_dice:\n            print(f'更新最佳DICE：({best_dice:0.4f}---->{val_dice:0.4f})')\n            best_dice = val_dice\n            best_epoch = epoch\n            best_model_wts = copy.deepcopy(model.state_dict())\n            if save:\n                model_save_path = f'/kaggle/working/save_model/fold{fold}_epoch_{best_epoch}.bin'\n                torch.save(best_model_wts, model_save_path)\n                print('模型已保存')\n\n        print('')\n        print('')\n    \n    if save:\n        model_save_path = f'/kaggle/working/save_model/fold{fold}_best_model.bin'\n        torch.save(best_model_wts, model_save_path)\n        print('最佳模型已保存')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 模型的训练循环","metadata":{}},{"cell_type":"code","source":"# 连接wandb平台\nwandb.login(key='b25ad7a07f1595e9696d4aa4df7a18fcd5881d90')\nwandb.init(project=\"Sartorius_unet\", entity=\"exiaoseiei\")\n\n# # 创建输出路径\n# os.makedirs('/kaggle/working/save_model')\n\n# 获取训练dataframe\ndataframe = get_df()\n\n# 暂时不采用交叉实验\ntrain_dataloader, valid_dataloader = get_fold_dataloader(fold=0, dataframe=dataframe)\nmodel = build_model_my()\nopt = torch.optim.Adam(model.parameters(), lr=cfg.lr, weight_decay=cfg.wd)\nscheduler = get_scheduler(opt)\nrun_epochs(model=model, optimizer=opt, scheduler=scheduler,\n           train_dataloader=train_dataloader, valid_dataloader=valid_dataloader, num_epoch=cfg.epochs, fold=0)","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}