Upload 15 files

TranSalNet_Res.py

@@ -0,0 +1,171 @@
+import os
+import torch
+import numpy as np
+import pandas as pd
+from torch.utils.data import Dataset, DataLoader
+from skimage import io, transform
+from PIL import Image
+import torch.nn as nn
+from torchvision import transforms, utils, models
+import torch.nn.functional as F
+import utils.resnet as resnet
+
+from utils.TransformerEncoder import Encoder
+
+
+
+cfg1 = {
+"hidden_size" : 768,
+"mlp_dim" : 768*4,
+"num_heads" : 12,
+"num_layers" : 2,
+"attention_dropout_rate" : 0,
+"dropout_rate" : 0.0,
+}
+
+cfg2 = {
+"hidden_size" : 768,
+"mlp_dim" : 768*4,
+"num_heads" : 12,
+"num_layers" : 2,
+"attention_dropout_rate" : 0,
+"dropout_rate" : 0.0,
+}
+
+cfg3 = {
+"hidden_size" : 512,
+"mlp_dim" : 512*4,
+"num_heads" : 8,
+"num_layers" : 2,
+"attention_dropout_rate" : 0,
+"dropout_rate" : 0.0,
+}
+
+
+class TranSalNet(nn.Module):
+
+    def __init__(self):
+        super(TranSalNet, self).__init__()
+        self.encoder = _Encoder()
+        self.decoder = _Decoder()
+
+    def forward(self, x):
+        x = self.encoder(x)
+        x = self.decoder(x)
+        return x
+
+
+class _Encoder(nn.Module):
+    def __init__(self):
+        super(_Encoder, self).__init__()
+        base_model = resnet.resnet50(pretrained=True)
+        base_layers = list(base_model.children())[:8]
+        self.encoder = nn.ModuleList(base_layers).eval()
+
+    def forward(self, x):
+        outputs = []
+        for ii,layer in enumerate(self.encoder):
+            x = layer(x)
+            if ii in {5,6,7}:
+                outputs.append(x)
+        return outputs
+
+
+class _Decoder(nn.Module):
+
+    def __init__(self):
+        super(_Decoder, self).__init__()
+        self.conv1 = nn.Conv2d(768, 768, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.conv2 = nn.Conv2d(768, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.conv3 = nn.Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.conv4 = nn.Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.conv5 = nn.Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.conv6 = nn.Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.conv7 = nn.Conv2d(32, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+
+        self.batchnorm1 = nn.BatchNorm2d(768, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+        self.batchnorm2 = nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+        self.batchnorm3 = nn.BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+        self.batchnorm4 = nn.BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+        self.batchnorm5 = nn.BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+        self.batchnorm6 = nn.BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+
+        self.TransEncoder1 = TransEncoder(in_channels=2048, spatial_size=9*12, cfg=cfg1)
+        self.TransEncoder2 = TransEncoder(in_channels=1024, spatial_size=18*24, cfg=cfg2)
+        self.TransEncoder3 = TransEncoder(in_channels=512, spatial_size=36*48, cfg=cfg3)
+
+        self.add = torch.add
+        self.relu = nn.ReLU(True)
+        self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        x3, x4, x5 = x
+
+        x5 = self.TransEncoder1(x5)
+        x5 = self.conv1(x5)
+        x5 = self.batchnorm1(x5)
+        x5 = self.relu(x5)
+        x5 = self.upsample(x5)
+
+        x4_a = self.TransEncoder2(x4)
+        x4 = x5 * x4_a
+        x4 = self.relu(x4)
+        x4 = self.conv2(x4)
+        x4 = self.batchnorm2(x4)
+        x4 = self.relu(x4)
+        x4 = self.upsample(x4)
+
+        x3_a = self.TransEncoder3(x3)
+        x3 = x4 * x3_a
+        x3 = self.relu(x3)
+        x3 = self.conv3(x3)
+        x3 = self.batchnorm3(x3)
+        x3 = self.relu(x3)
+        x3 = self.upsample(x3)
+
+        x2 = self.conv4(x3)
+        x2 = self.batchnorm4(x2)
+        x2 = self.relu(x2)
+        x2 = self.upsample(x2)
+        x2 = self.conv5(x2)
+        x2 = self.batchnorm5(x2)
+        x2 = self.relu(x2)
+
+        x1 = self.upsample(x2)
+        x1 = self.conv6(x1)
+        x1 = self.batchnorm6(x1)
+        x1 = self.relu(x1)
+        x1 = self.conv7(x1)
+        x = self.sigmoid(x1)
+
+        return x
+
+
+class TransEncoder(nn.Module):
+
+    def __init__(self, in_channels, spatial_size, cfg):
+        super(TransEncoder, self).__init__()
+
+        self.patch_embeddings = nn.Conv2d(in_channels=in_channels,
+                                          out_channels=cfg['hidden_size'],
+                                          kernel_size=1,
+                                          stride=1)
+        self.position_embeddings = nn.Parameter(torch.zeros(1, spatial_size, cfg['hidden_size']))
+
+        self.transformer_encoder = Encoder(cfg)
+
+    def forward(self, x):
+        a, b = x.shape[2], x.shape[3]
+        x = self.patch_embeddings(x)
+        x = x.flatten(2)
+        x = x.transpose(-1, -2)
+
+        embeddings = x + self.position_embeddings
+        x = self.transformer_encoder(embeddings)
+        B, n_patch, hidden = x.shape
+        x = x.permute(0, 2, 1)
+        x = x.contiguous().view(B, hidden, a, b)
+
+        return x
+

app.py

@@ -0,0 +1,54 @@
+import streamlit as st
+import torch
+import cv2
+from PIL import Image
+import numpy as np
+from torchvision import transforms
+from TranSalNet_Res import TranSalNet  # Make sure TranSalNet is accessible from your Streamlit app
+
+# Load the model and set the device
+model = TranSalNet()
+model.load_state_dict(torch.load('pretrained_models/TranSalNet_Res.pth', map_location=torch.device('cpu')))
+model.eval()  # Set the model to evaluation mode
+device = torch.device('cpu')
+model.to(device)
+
+# Define Streamlit app
+st.title('Saliency Detection App')
+st.write('Upload an image for saliency detection:')
+uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+
+if uploaded_image:
+    image = Image.open(uploaded_image)
+    st.image(image, caption='Uploaded Image', use_column_width=True)
+
+    # Check if the user clicks a button
+    if st.button('Detect Saliency'):
+        # Preprocess the image
+        img = image.resize((384, 288))
+        img = np.array(img) / 255.
+        img = np.transpose(img, (2, 0, 1))
+        img = torch.from_numpy(img).unsqueeze(0).float()
+        img = img.to(device)
+
+        # Get saliency prediction
+        with torch.no_grad():
+            pred_saliency = model(img)
+
+        # Convert the result back to a PIL image
+        toPIL = transforms.ToPILImage()
+        pic = toPIL(pred_saliency.squeeze())
+
+        # Colorize the grayscale prediction
+        colorized_img = cv2.applyColorMap(np.uint8(pic), cv2.COLORMAP_JET)
+
+        # Ensure the colorized image has the same dimensions as the original image
+        original_img = np.array(image)
+        colorized_img = cv2.resize(colorized_img, (original_img.shape[1], original_img.shape[0]))
+
+        # You can add more post-processing here if needed
+
+        # Display the final result
+        st.image(colorized_img, caption='Colorized Saliency Map', use_column_width=True)
+
+st.write('Finished!')

pretrained_models/.keep

@@ -0,0 +1 @@
+

pretrained_models/TranSalNet_Res.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3853e24e1e0bf892bdc321dcf269516a58450f9af2d3ca0b620272d6c81fe5c7
+size 290451767

pretrained_models/resnet50-0676ba61.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0676ba61b6795bbe1773cffd859882e5e297624d384b6993f7c9e683e722fb8a
+size 102530333

utils/TransformerEncoder.py

@@ -0,0 +1,137 @@
+# coding=utf-8
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import copy
+import logging
+import math
+
+from os.path import join as pjoin
+
+import torch
+import torch.nn as nn
+import numpy as np
+
+from torch.nn import CrossEntropyLoss, Dropout, Softmax, Linear, Conv2d, LayerNorm
+from torch.nn.modules.utils import _pair
+from scipy import ndimage
+
+
+ACT2FN = {"gelu": torch.nn.functional.gelu, "relu": torch.nn.functional.relu}
+
+
+class Attention(nn.Module):
+    def __init__(self, config):
+        super(Attention, self).__init__()
+        self.num_attention_heads = config["num_heads"]  # 12
+        self.attention_head_size = int(config['hidden_size'] / self.num_attention_heads)    # 42
+        self.all_head_size = self.num_attention_heads * self.attention_head_size    # 12*42=504
+
+        self.query = Linear(config['hidden_size'], self.all_head_size)  # (512, 504)
+        self.key = Linear(config['hidden_size'], self.all_head_size)
+        self.value = Linear(config['hidden_size'], self.all_head_size)
+
+        # self.out = Linear(config['hidden_size'], config['hidden_size'])
+        self.out = Linear(self.all_head_size, config['hidden_size'])
+        self.attn_dropout = Dropout(config["attention_dropout_rate"])
+        self.proj_dropout = Dropout(config["attention_dropout_rate"])
+
+        self.softmax = Softmax(dim=-1)
+
+    def transpose_for_scores(self, x):
+        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+        x = x.view(*new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward(self, hidden_states):
+
+        mixed_query_layer = self.query(hidden_states)
+        mixed_key_layer = self.key(hidden_states)
+        mixed_value_layer = self.value(hidden_states)
+
+        query_layer = self.transpose_for_scores(mixed_query_layer)
+        key_layer = self.transpose_for_scores(mixed_key_layer)
+        value_layer = self.transpose_for_scores(mixed_value_layer)
+
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+        attention_probs = self.softmax(attention_scores)
+        attention_probs = self.attn_dropout(attention_probs)
+
+        context_layer = torch.matmul(attention_probs, value_layer)
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(*new_context_layer_shape)
+        attention_output = self.out(context_layer)
+        attention_output = self.proj_dropout(attention_output)
+        return attention_output
+
+
+class Mlp(nn.Module):
+    def __init__(self, config):
+        super(Mlp, self).__init__()
+        self.fc1 = Linear(config['hidden_size'], config["mlp_dim"])
+        self.fc2 = Linear(config["mlp_dim"], config['hidden_size'])
+        self.act_fn = ACT2FN["gelu"]
+        self.dropout = Dropout(config["dropout_rate"])
+        self._init_weights()
+
+    def _init_weights(self):
+        nn.init.xavier_uniform_(self.fc1.weight)
+        nn.init.xavier_uniform_(self.fc2.weight)
+        nn.init.normal_(self.fc1.bias, std=1e-6)
+        nn.init.normal_(self.fc2.bias, std=1e-6)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act_fn(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+        x = self.dropout(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, config):
+        super(Block, self).__init__()
+        self.flag = config['num_heads']
+        self.hidden_size = config['hidden_size']
+        self.ffn_norm = LayerNorm(config['hidden_size'], eps=1e-6)
+        self.ffn = Mlp(config)
+        self.attn = Attention(config)
+        self.attention_norm = LayerNorm(config['hidden_size'], eps=1e-6)
+
+    def forward(self, x):
+        h = x
+
+        x = self.attention_norm(x)
+        x = self.attn(x)
+        x = x + h
+
+        h = x
+        x = self.ffn_norm(x)
+        x = self.ffn(x)
+        x = x + h
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(self, config):
+        super(Encoder, self).__init__()
+
+        self.layer = nn.ModuleList()
+        self.encoder_norm = LayerNorm(config['hidden_size'], eps=1e-6)
+        for _ in range(config["num_layers"]):
+            layer = Block(config)
+            self.layer.append(copy.deepcopy(layer))
+
+    def forward(self, hidden_states):
+        for layer_block in self.layer:
+            hidden_states = layer_block(hidden_states)
+        encoded = self.encoder_norm(hidden_states)
+
+        return encoded
+
+

utils/data_process.py

@@ -0,0 +1,116 @@
+import cv2
+from PIL import Image
+import numpy as np
+import os
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+
+def preprocess_img(img_dir, channels=3):
+
+    if channels == 1:
+        img = cv2.imread(img_dir, 0)
+    elif channels == 3:
+        img = cv2.imread(img_dir)
+
+    shape_r = 288
+    shape_c = 384
+    img_padded = np.ones((shape_r, shape_c, channels), dtype=np.uint8)
+    if channels == 1:
+        img_padded = np.zeros((shape_r, shape_c), dtype=np.uint8)
+    original_shape = img.shape
+    rows_rate = original_shape[0] / shape_r
+    cols_rate = original_shape[1] / shape_c
+    if rows_rate > cols_rate:
+        new_cols = (original_shape[1] * shape_r) // original_shape[0]
+        img = cv2.resize(img, (new_cols, shape_r))
+        if new_cols > shape_c:
+            new_cols = shape_c
+        img_padded[:,
+        ((img_padded.shape[1] - new_cols) // 2):((img_padded.shape[1] - new_cols) // 2 + new_cols)] = img
+    else:
+        new_rows = (original_shape[0] * shape_c) // original_shape[1]
+        img = cv2.resize(img, (shape_c, new_rows))
+
+        if new_rows > shape_r:
+            new_rows = shape_r
+        img_padded[((img_padded.shape[0] - new_rows) // 2):((img_padded.shape[0] - new_rows) // 2 + new_rows),
+        :] = img
+
+    return img_padded
+
+
+def postprocess_img(pred, org_dir):
+    pred = np.array(pred)
+    org = cv2.imread(org_dir, 0)
+    shape_r = org.shape[0]
+    shape_c = org.shape[1]
+    predictions_shape = pred.shape
+
+    rows_rate = shape_r / predictions_shape[0]
+    cols_rate = shape_c / predictions_shape[1]
+
+    if rows_rate > cols_rate:
+        new_cols = (predictions_shape[1] * shape_r) // predictions_shape[0]
+        pred = cv2.resize(pred, (new_cols, shape_r))
+        img = pred[:, ((pred.shape[1] - shape_c) // 2):((pred.shape[1] - shape_c) // 2 + shape_c)]
+    else:
+        new_rows = (predictions_shape[0] * shape_c) // predictions_shape[1]
+        pred = cv2.resize(pred, (shape_c, new_rows))
+        img = pred[((pred.shape[0] - shape_r) // 2):((pred.shape[0] - shape_r) // 2 + shape_r), :]
+
+    return img
+
+
+class MyDataset(Dataset):
+    """Load dataset."""
+
+    def __init__(self, ids, stimuli_dir, saliency_dir, fixation_dir, transform=None):
+        """
+        Args:
+            csv_file (string): Path to the csv file with annotations.
+            root_dir (string): Directory with all the images.
+            transform (callable, optional): Optional transform to be applied
+                on a sample.
+        """
+        self.ids = ids
+        self.stimuli_dir = stimuli_dir
+        self.saliency_dir = saliency_dir
+        self.fixation_dir = fixation_dir
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.ids)
+
+    def __getitem__(self, idx):
+        if torch.is_tensor(idx):
+            idx = idx.tolist()
+
+        im_path = self.stimuli_dir + self.ids.iloc[idx, 0]
+        image = Image.open(im_path).convert('RGB')
+        img = np.array(image) / 255.
+        img = np.transpose(img, (2, 0, 1))
+        img = torch.from_numpy(img)
+        # if self.transform:
+        #    img = self.transform(image)
+
+        smap_path = self.saliency_dir + self.ids.iloc[idx, 1]
+        saliency = Image.open(smap_path)
+
+        smap = np.expand_dims(np.array(saliency) / 255., axis=0)
+        smap = torch.from_numpy(smap)
+
+        fmap_path = self.fixation_dir + self.ids.iloc[idx, 2]
+        fixation = Image.open(fmap_path)
+
+        fmap = np.expand_dims(np.array(fixation) / 255., axis=0)
+        fmap = torch.from_numpy(fmap)
+
+        sample = {'image': img, 'saliency': smap, 'fixation': fmap}
+
+        return sample
+
+
+
+
+

utils/densenet.py

@@ -0,0 +1,287 @@
+import re
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as cp
+from collections import OrderedDict
+# from .utils import load_state_dict_from_url
+from torch import Tensor
+from torch.jit.annotations import List
+
+
+__all__ = ['DenseNet', 'densenet121', 'densenet169', 'densenet201', 'densenet161']
+
+model_urls = {
+    'densenet121': 'https://download.pytorch.org/models/densenet121-a639ec97.pth',
+    'densenet169': 'https://download.pytorch.org/models/densenet169-b2777c0a.pth',
+    'densenet201': 'https://download.pytorch.org/models/densenet201-c1103571.pth',
+    'densenet161': 'https://download.pytorch.org/models/densenet161-8d451a50.pth',
+}
+
+
+class _DenseLayer(nn.Module):
+    def __init__(self, num_input_features, growth_rate, bn_size, drop_rate, memory_efficient=False):
+        super(_DenseLayer, self).__init__()
+        self.add_module('norm1', nn.BatchNorm2d(num_input_features)),
+        self.add_module('relu1', nn.ReLU(inplace=True)),
+        self.add_module('conv1', nn.Conv2d(num_input_features, bn_size *
+                                           growth_rate, kernel_size=1, stride=1,
+                                           bias=False)),
+        self.add_module('norm2', nn.BatchNorm2d(bn_size * growth_rate)),
+        self.add_module('relu2', nn.ReLU(inplace=True)),
+        self.add_module('conv2', nn.Conv2d(bn_size * growth_rate, growth_rate,
+                                           kernel_size=3, stride=1, padding=1,
+                                           bias=False)),
+        self.drop_rate = float(drop_rate)
+        self.memory_efficient = memory_efficient
+
+    def bn_function(self, inputs):
+        # type: (List[Tensor]) -> Tensor
+        concated_features = torch.cat(inputs, 1)
+        bottleneck_output = self.conv1(self.relu1(self.norm1(concated_features)))  # noqa: T484
+        return bottleneck_output
+
+    # todo: rewrite when torchscript supports any
+    def any_requires_grad(self, input):
+        # type: (List[Tensor]) -> bool
+        for tensor in input:
+            if tensor.requires_grad:
+                return True
+        return False
+
+    @torch.jit.unused  # noqa: T484
+    def call_checkpoint_bottleneck(self, input):
+        # type: (List[Tensor]) -> Tensor
+        def closure(*inputs):
+            return self.bn_function(inputs)
+
+        return cp.checkpoint(closure, *input)
+
+    @torch.jit._overload_method  # noqa: F811
+    def forward(self, input):
+        # type: (List[Tensor]) -> (Tensor)
+        pass
+
+    @torch.jit._overload_method  # noqa: F811
+    def forward(self, input):
+        # type: (Tensor) -> (Tensor)
+        pass
+
+    # torchscript does not yet support *args, so we overload method
+    # allowing it to take either a List[Tensor] or single Tensor
+    def forward(self, input):  # noqa: F811
+        if isinstance(input, Tensor):
+            prev_features = [input]
+        else:
+            prev_features = input
+
+        if self.memory_efficient and self.any_requires_grad(prev_features):
+            if torch.jit.is_scripting():
+                raise Exception("Memory Efficient not supported in JIT")
+
+            bottleneck_output = self.call_checkpoint_bottleneck(prev_features)
+        else:
+            bottleneck_output = self.bn_function(prev_features)
+
+        new_features = self.conv2(self.relu2(self.norm2(bottleneck_output)))
+        if self.drop_rate > 0:
+            new_features = F.dropout(new_features, p=self.drop_rate,
+                                     training=self.training)
+        return new_features
+
+
+class _DenseBlock(nn.ModuleDict):
+    _version = 2
+
+    def __init__(self, num_layers, num_input_features, bn_size, growth_rate, drop_rate, memory_efficient=False):
+        super(_DenseBlock, self).__init__()
+        for i in range(num_layers):
+            layer = _DenseLayer(
+                num_input_features + i * growth_rate,
+                growth_rate=growth_rate,
+                bn_size=bn_size,
+                drop_rate=drop_rate,
+                memory_efficient=memory_efficient,
+            )
+            self.add_module('denselayer%d' % (i + 1), layer)
+
+    def forward(self, init_features):
+        features = [init_features]
+        for name, layer in self.items():
+            new_features = layer(features)
+            features.append(new_features)
+        return torch.cat(features, 1)
+
+
+class _Transition(nn.Sequential):
+    def __init__(self, num_input_features, num_output_features):
+        super(_Transition, self).__init__()
+        self.add_module('norm', nn.BatchNorm2d(num_input_features))
+        self.add_module('relu', nn.ReLU(inplace=True))
+        self.add_module('conv', nn.Conv2d(num_input_features, num_output_features,
+                                          kernel_size=1, stride=1, bias=False))
+        self.add_module('pool', nn.AvgPool2d(kernel_size=2, stride=2))
+
+
+class DenseNet(nn.Module):
+    r"""Densenet-BC model class, based on
+    `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_
+
+    Args:
+        growth_rate (int) - how many filters to add each layer (`k` in paper)
+        block_config (list of 4 ints) - how many layers in each pooling block
+        num_init_features (int) - the number of filters to learn in the first convolution layer
+        bn_size (int) - multiplicative factor for number of bottle neck layers
+          (i.e. bn_size * k features in the bottleneck layer)
+        drop_rate (float) - dropout rate after each dense layer
+        num_classes (int) - number of classification classes
+        memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient,
+          but slower. Default: *False*. See `"paper" <https://arxiv.org/pdf/1707.06990.pdf>`_
+    """
+
+    def __init__(self, growth_rate=32, block_config=(6, 12, 24, 16),
+                 num_init_features=64, bn_size=4, drop_rate=0, num_classes=1000, memory_efficient=False):
+
+        super(DenseNet, self).__init__()
+
+        # First convolution
+        self.features = nn.Sequential(OrderedDict([
+            ('conv0', nn.Conv2d(3, num_init_features, kernel_size=7, stride=2,
+                                padding=3, bias=False)),
+            ('norm0', nn.BatchNorm2d(num_init_features)),
+            ('relu0', nn.ReLU(inplace=True)),
+            ('pool0', nn.MaxPool2d(kernel_size=3, stride=2, padding=1)),
+        ]))
+
+        # Each denseblock
+        num_features = num_init_features
+        for i, num_layers in enumerate(block_config):
+            block = _DenseBlock(
+                num_layers=num_layers,
+                num_input_features=num_features,
+                bn_size=bn_size,
+                growth_rate=growth_rate,
+                drop_rate=drop_rate,
+                memory_efficient=memory_efficient
+            )
+            self.features.add_module('denseblock%d' % (i + 1), block)
+            num_features = num_features + num_layers * growth_rate
+            if i != len(block_config) - 1:
+                trans = _Transition(num_input_features=num_features,
+                                    num_output_features=num_features // 2)
+                self.features.add_module('transition%d' % (i + 1), trans)
+                num_features = num_features // 2
+
+        # Final batch norm
+        self.features.add_module('norm5', nn.BatchNorm2d(num_features))
+
+        # Linear layer
+        self.classifier = nn.Linear(num_features, num_classes)
+
+        # Official init from torch repo.
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        features = self.features(x)
+        out = F.relu(features, inplace=True)
+        out = F.adaptive_avg_pool2d(out, (1, 1))
+        out = torch.flatten(out, 1)
+        out = self.classifier(out)
+        return out
+
+
+def _load_state_dict(model, model_url, progress, flag):
+    # '.'s are no longer allowed in module names, but previous _DenseLayer
+    # has keys 'norm.1', 'relu.1', 'conv.1', 'norm.2', 'relu.2', 'conv.2'.
+    # They are also in the checkpoints in model_urls. This pattern is used
+    # to find such keys.
+    pattern = re.compile(
+        r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$')
+    if flag == "densenet161":
+        state_dict = torch.load(r'pretrained_models/densenet161-8d451a50.pth')
+    else:
+        state_dict = load_state_dict_from_url(model_url, progress=progress)
+    for key in list(state_dict.keys()):
+        res = pattern.match(key)
+        if res:
+            new_key = res.group(1) + res.group(2)
+            state_dict[new_key] = state_dict[key]
+            del state_dict[key]
+    model.load_state_dict(state_dict)
+
+
+def _densenet(arch, growth_rate, block_config, num_init_features, pretrained, progress,
+              **kwargs):
+    model = DenseNet(growth_rate, block_config, num_init_features, **kwargs)
+    if pretrained:
+        if arch == 'densenet161':
+            _load_state_dict(model, model_urls[arch], progress, 'densenet161')
+        else:
+            _load_state_dict(model, model_urls[arch], progress, 0)
+    return model
+
+
+def densenet121(pretrained=False, progress=True, **kwargs):
+    r"""Densenet-121 model from
+    `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+        memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient,
+          but slower. Default: *False*. See `"paper" <https://arxiv.org/pdf/1707.06990.pdf>`_
+    """
+    return _densenet('densenet121', 32, (6, 12, 24, 16), 64, pretrained, progress,
+                     **kwargs)
+
+
+
+def densenet161(pretrained=False, progress=True, **kwargs):
+    r"""Densenet-161 model from
+    `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+        memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient,
+          but slower. Default: *False*. See `"paper" <https://arxiv.org/pdf/1707.06990.pdf>`_
+    """
+    return _densenet('densenet161', 48, (6, 12, 36, 24), 96, pretrained, progress,
+                     **kwargs)
+
+
+
+def densenet169(pretrained=False, progress=True, **kwargs):
+    r"""Densenet-169 model from
+    `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+        memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient,
+          but slower. Default: *False*. See `"paper" <https://arxiv.org/pdf/1707.06990.pdf>`_
+    """
+    return _densenet('densenet169', 32, (6, 12, 32, 32), 64, pretrained, progress,
+                     **kwargs)
+
+
+
+def densenet201(pretrained=False, progress=True, **kwargs):
+    r"""Densenet-201 model from
+    `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+        memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient,
+          but slower. Default: *False*. See `"paper" <https://arxiv.org/pdf/1707.06990.pdf>`_
+    """
+    return _densenet('densenet201', 32, (6, 12, 48, 32), 64, pretrained, progress,
+                     **kwargs)

utils/loss_function.py

@@ -0,0 +1,69 @@
+import torch as t
+import torch.nn as nn
+#import numpy as np
+
+
+class SaliencyLoss(nn.Module):
+    def __init__(self):
+        super(SaliencyLoss, self).__init__()
+
+    def forward(self, preds, labels, loss_type='cc'):
+        losses = []
+        if loss_type == 'cc':
+            for i in range(labels.shape[0]): # labels.shape[0] is batch size
+                loss = loss_CC(preds[i],labels[i])
+                losses.append(loss)
+
+        elif loss_type == 'kldiv':
+            for i in range(labels.shape[0]):
+                loss = loss_KLdiv(preds[i],labels[i])
+                losses.append(loss)
+
+        elif loss_type == 'sim':
+            for i in range(labels.shape[0]):
+                loss = loss_similarity(preds[i],labels[i])
+                losses.append(loss)
+
+        elif loss_type == 'nss':
+            for i in range(labels.shape[0]):
+                loss = loss_NSS(preds[i],labels[i])
+                losses.append(loss)
+            
+        return t.stack(losses).mean(dim=0, keepdim=True)
+        
+        
+def loss_KLdiv(pred_map, gt_map):
+    eps = 2.2204e-16
+    pred_map = pred_map/t.sum(pred_map)
+    gt_map = gt_map/t.sum(gt_map)
+    div = t.sum(t.mul(gt_map, t.log(eps + t.div(gt_map,pred_map+eps))))
+    return div 
+        
+    
+def loss_CC(pred_map,gt_map):
+    gt_map_ = (gt_map - t.mean(gt_map))
+    pred_map_ = (pred_map - t.mean(pred_map))
+    cc = t.sum(t.mul(gt_map_,pred_map_))/t.sqrt(t.sum(t.mul(gt_map_,gt_map_))*t.sum(t.mul(pred_map_,pred_map_)))
+    return cc
+
+
+def loss_similarity(pred_map,gt_map):
+    gt_map = (gt_map - t.min(gt_map))/(t.max(gt_map)-t.min(gt_map))
+    gt_map = gt_map/t.sum(gt_map)
+    
+    pred_map = (pred_map - t.min(pred_map))/(t.max(pred_map)-t.min(pred_map))
+    pred_map = pred_map/t.sum(pred_map)
+    
+    diff = t.min(gt_map,pred_map)
+    score = t.sum(diff)
+    
+    return score
+    
+    
+def loss_NSS(pred_map,fix_map):
+    '''ground truth here is fixation map'''
+
+    pred_map_ = (pred_map - t.mean(pred_map))/t.std(pred_map)
+    mask = fix_map.gt(0)
+    score = t.mean(t.masked_select(pred_map_, mask))
+    return score

utils/resnet.py

@@ -0,0 +1,419 @@
+from typing import Type, Any, Callable, Union, List, Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+# from .._internally_replaced_utils import load_state_dict_from_url
+# from ..utils import _log_api_usage_once
+
+
+__all__ = [
+    "ResNet",
+    "resnet18",
+    "resnet34",
+    "resnet50",
+    "resnet101",
+    "resnet152",
+    "resnext50_32x4d",
+    "resnext101_32x8d",
+    "wide_resnet50_2",
+    "wide_resnet101_2",
+]
+
+
+model_urls = {
+    "resnet18": "https://download.pytorch.org/models/resnet18-f37072fd.pth",
+    "resnet34": "https://download.pytorch.org/models/resnet34-b627a593.pth",
+    "resnet50": "https://download.pytorch.org/models/resnet50-0676ba61.pth",
+    "resnet101": "https://download.pytorch.org/models/resnet101-63fe2227.pth",
+    "resnet152": "https://download.pytorch.org/models/resnet152-394f9c45.pth",
+    "resnext50_32x4d": "https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth",
+    "resnext101_32x8d": "https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth",
+    "wide_resnet50_2": "https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth",
+    "wide_resnet101_2": "https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth",
+}
+
+
+def conv3x3(in_planes: int, out_planes: int, stride: int = 1, groups: int = 1, dilation: int = 1) -> nn.Conv2d:
+    """3x3 convolution with padding"""
+    return nn.Conv2d(
+        in_planes,
+        out_planes,
+        kernel_size=3,
+        stride=stride,
+        padding=dilation,
+        groups=groups,
+        bias=False,
+        dilation=dilation,
+    )
+
+
+def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion: int = 1
+
+    def __init__(
+        self,
+        inplanes: int,
+        planes: int,
+        stride: int = 1,
+        downsample: Optional[nn.Module] = None,
+        groups: int = 1,
+        base_width: int = 64,
+        dilation: int = 1,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+    ) -> None:
+        super().__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError("BasicBlock only supports groups=1 and base_width=64")
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
+    # This variant is also known as ResNet V1.5 and improves accuracy according to
+    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+
+    expansion: int = 4
+
+    def __init__(
+        self,
+        inplanes: int,
+        planes: int,
+        stride: int = 1,
+        downsample: Optional[nn.Module] = None,
+        groups: int = 1,
+        base_width: int = 64,
+        dilation: int = 1,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+    ) -> None:
+        super().__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(planes * (base_width / 64.0)) * groups
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(inplanes, width)
+        self.bn1 = norm_layer(width)
+        self.conv2 = conv3x3(width, width, stride, groups, dilation)
+        self.bn2 = norm_layer(width)
+        self.conv3 = conv1x1(width, planes * self.expansion)
+        self.bn3 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(
+        self,
+        block: Type[Union[BasicBlock, Bottleneck]],
+        layers: List[int],
+        num_classes: int = 1000,
+        zero_init_residual: bool = False,
+        groups: int = 1,
+        width_per_group: int = 64,
+        replace_stride_with_dilation: Optional[List[bool]] = None,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+    ) -> None:
+        super().__init__()
+        # _log_api_usage_once(self)
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+
+        self.inplanes = 64
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            # each element in the tuple indicates if we should replace
+            # the 2x2 stride with a dilated convolution instead
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError(
+                "replace_stride_with_dilation should be None "
+                f"or a 3-element tuple, got {replace_stride_with_dilation}"
+            )
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2, dilate=replace_stride_with_dilation[0])
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1])
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2])
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)  # type: ignore[arg-type]
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)  # type: ignore[arg-type]
+
+    def _make_layer(
+        self,
+        block: Type[Union[BasicBlock, Bottleneck]],
+        planes: int,
+        blocks: int,
+        stride: int = 1,
+        dilate: bool = False,
+    ) -> nn.Sequential:
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(
+            block(
+                self.inplanes, planes, stride, downsample, self.groups, self.base_width, previous_dilation, norm_layer
+            )
+        )
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(
+                block(
+                    self.inplanes,
+                    planes,
+                    groups=self.groups,
+                    base_width=self.base_width,
+                    dilation=self.dilation,
+                    norm_layer=norm_layer,
+                )
+            )
+
+        return nn.Sequential(*layers)
+
+    def _forward_impl(self, x: Tensor) -> Tensor:
+        # See note [TorchScript super()]
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+
+        return x
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self._forward_impl(x)
+
+
+def _resnet(
+    arch: str,
+    block: Type[Union[BasicBlock, Bottleneck]],
+    layers: List[int],
+    pretrained: bool,
+    progress: bool,
+    **kwargs: Any,
+) -> ResNet:
+    model = ResNet(block, layers, **kwargs)
+    if pretrained:
+        if arch == 'resnet50':
+            state_dict = torch.load(r'pretrained_models/resnet50-0676ba61.pth')
+        else:
+            state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)
+        model.load_state_dict(state_dict)
+    return model
+
+
+def resnet18(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+    r"""ResNet-18 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet("resnet18", BasicBlock, [2, 2, 2, 2], pretrained, progress, **kwargs)
+
+
+def resnet34(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+    r"""ResNet-34 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet("resnet34", BasicBlock, [3, 4, 6, 3], pretrained, progress, **kwargs)
+
+
+def resnet50(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+    r"""ResNet-50 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet("resnet50", Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs)
+
+
+def resnet101(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+    r"""ResNet-101 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet("resnet101", Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs)
+
+
+def resnet152(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+    r"""ResNet-152 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet("resnet152", Bottleneck, [3, 8, 36, 3], pretrained, progress, **kwargs)
+
+
+def resnext50_32x4d(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+    r"""ResNeXt-50 32x4d model from
+    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs["groups"] = 32
+    kwargs["width_per_group"] = 4
+    return _resnet("resnext50_32x4d", Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs)
+
+
+def resnext101_32x8d(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+    r"""ResNeXt-101 32x8d model from
+    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs["groups"] = 32
+    kwargs["width_per_group"] = 8
+    return _resnet("resnext101_32x8d", Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs)
+
+
+
+def wide_resnet50_2(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+    r"""Wide ResNet-50-2 model from
+    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_.
+
+    The model is the same as ResNet except for the bottleneck number of channels
+    which is twice larger in every block. The number of channels in outer 1x1
+    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
+    channels, and in Wide ResNet-50-2 has 2048-1024-2048.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs["width_per_group"] = 64 * 2
+    return _resnet("wide_resnet50_2", Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs)
+
+
+
+
+def wide_resnet101_2(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+    r"""Wide ResNet-101-2 model from
+    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_.
+
+    The model is the same as ResNet except for the bottleneck number of channels
+    which is twice larger in every block. The number of channels in outer 1x1
+    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
+    channels, and in Wide ResNet-50-2 has 2048-1024-2048.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs["width_per_group"] = 64 * 2
+    return _resnet("wide_resnet101_2", Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs)