{"cells":[{"metadata":{},"cell_type":"markdown","source":"# ResNet200d encoder for segmentation_models_pytorch\n\nThere isn't resnet200d encoder included in segmentation_models_pytorch package.\nI'm going to show a simple way to make your own encoder using pytorch_image_models package."},{"metadata":{},"cell_type":"markdown","source":"## 1. Install relevant packages"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":false,"_kg_hide-output":true,"collapsed":true},"cell_type":"code","source":"!pip install timm\n!pip install segmentation_models_pytorch","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 2. Define encoder"},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom timm import create_model\nimport segmentation_models_pytorch as smp\nfrom segmentation_models_pytorch.encoders._base import EncoderMixin\n\n\nclass ResNet200dEncoder(nn.Module, EncoderMixin):\n    def __init__(self, **kwargs):\n        super().__init__()\n        self._out_channels = [3, 64, 256, 512, 1024, 2048] # output channels\n        self._depth = 5 # UNet depth\n        self._in_channels = 3\n        self._m = create_model(\n            model_name='resnet200d',\n            in_chans=self._in_channels,\n            pretrained=False\n        )\n        self._m.global_pool = nn.Identity()\n        self._m.fc = nn.Identity()\n\n    def get_stages(self):\n        return [\n            nn.Identity(),\n            nn.Sequential(self._m.conv1, self._m.bn1, self._m.act1),\n            nn.Sequential(self._m.maxpool, self._m.layer1),\n            self._m.layer2,\n            self._m.layer3,\n            self._m.layer4,\n        ]\n\n    def forward(self, x):\n        stages = self.get_stages()\n        features = []\n        for i in range(self._depth + 1):\n            x = stages[i](x)\n            features.append(x)\n        return features\n\n    def load_state_dict(self, state_dict):\n        state_dict.pop(\"fc.bias\")\n        state_dict.pop(\"fc.weight\")\n        self._m.load_state_dict(state_dict, strict=False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 3. Register new encoder and pretrained weight"},{"metadata":{"trusted":true},"cell_type":"code","source":"smp.encoders.encoders[\"resnet200d\"] = {\n    \"encoder\": ResNet200dEncoder,\n    \"pretrained_settings\": {\n        \"imagenet\": {\n            \"mean\": [0.485, 0.456, 0.406],\n            \"std\": [0.229, 0.224, 0.225],\n            \"url\": \"https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/resnet200d_ra2-bdba9bf9.pth\", # this can be the pretrained weight from timm\n            \"input_space\": \"RGB\",\n            \"input_range\": [0, 1],\n        },\n    },\n    \"params\": {}\n}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 4. Tadaa"},{"metadata":{"trusted":true},"cell_type":"code","source":"smp.Unet(\n    encoder_name=\"resnet200d\",        \n    encoder_weights=\"imagenet\", \n)","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}