vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/imatch-toolbox/third_party/patch2pix/networks/patch2pix.py

@@ -0,0 +1,403 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+import networks.resnet as resnet
+from networks.modules import *
+from networks.utils import select_local_patch_feats, filter_coarse
+from networks.ncn.model import MutualMatching, NeighConsensus  
+from networks.ncn.extract_ncmatches import corr_to_matches
+
+class Patch2Pix(nn.Module):    
+    def __init__(self, config):
+        super().__init__()
+        self.device = config.device
+        self.backbone = config.backbone
+        self.change_stride = config.change_stride        
+        self.upsample = 16
+        self.feats_downsample = [1, 2, 2, 2, 2]
+        feat_dims = [3, 64, 64, 128, 256]  # Resnet34 block feature out dims 
+        
+        # Initialize necessary network components
+        self.extract = resnet.__dict__[self.backbone]()
+        if self.change_stride:
+            self.extract.change_stride(target='layer3')
+            self.upsample //= 2            
+            self.feats_downsample[-1] = 1        
+        print(f'Initialize Patch2Pix: backbone={self.backbone} '
+              f'cstride={self.change_stride} upsample={self.upsample}')
+        
+        self.combine = FeatCorrelation(shape='4D')
+        self.ncn = NeighConsensus(kernel_sizes=[3, 3], channels=[16, 1])
+        
+        # Initialize regressor
+        self.regressor_config = config.regressor_config
+        if not self.regressor_config:
+            # If no regressor defined, model only computes coarse matches
+            self.regress_mid = None
+            self.regress_fine = None
+        else:
+            print(f'Init regressor {self.regressor_config}')
+            self.regr_batch = config.regr_batch
+            self.feat_idx = config.feat_idx            
+            feat_dim = 0  # Regressor's input feature dim
+            for idx in self.feat_idx:
+                feat_dim += feat_dims[idx]           
+            self.regressor_config.feat_dim = feat_dim            
+            self.ptype = ['center', 'center']
+            self.psize = config.regressor_config.psize            
+            self.pshift = config.regressor_config.pshift
+            self.panc = config.regressor_config.panc
+            self.shared = config.regressor_config.shared
+            self.regress_mid = FeatRegressNet(self.regressor_config, psize=self.psize[0])
+            if self.shared:
+                self.regress_fine = self.regress_mid
+                self.psize[1] = self.psize[0]
+            else:
+                self.regress_fine = FeatRegressNet(self.regressor_config, psize=self.psize[1])
+            
+        self.to(self.device)
+        self.init_weights_(weights_dict=config.weights_dict, pretrained=True)        
+           
+        if config.training:
+            self.freeze_feat = config.freeze_feat
+            # Freeze (part of) the backbone
+            print('Freezing feature extractor params upto layer {}'.format(self.freeze_feat))
+            for i, param in enumerate(self.extract.parameters()):         
+                # Resnet34 layer3=[48:87] blocks:
+                # 0=[48:57] 1=[57:63] 2=[63:69]
+                # 3=[69:75] 4=[75:81] 5=[81:87] 
+                if i < self.freeze_feat:
+                    param.requires_grad = False
+
+                # Always freeze resnet layer4, since never used
+                if i >= 87:
+                    param.requires_grad = False        
+            
+            config.optim_config.start_epoch = config.start_epoch
+            self.set_optimizer_(config.optim_config)
+    
+    def set_optimizer_(self, optim_config): 
+        params = []
+        if self.regress_mid:
+            params += list(self.regress_mid.parameters()) 
+        if self.regress_fine and not self.shared:
+            params += list(self.regress_fine.parameters())         
+        params += list(self.ncn.parameters())
+        if self.freeze_feat < 87:
+            params += list(self.extract.parameters())[self.freeze_feat:87]
+        self.optimizer, self.lr_scheduler = init_optimizer(params, optim_config)
+        print('Init optimizer, items: {}'.format(len(params)))
+
+    def optim_step_(self, loss):
+        self.optimizer.zero_grad()
+        loss.backward()
+        self.optimizer.step()            
+                
+    def init_weights_(self, weights_dict=None, pretrained=True):
+        print('Xavier initialize all model parameters')
+        self.apply(xavier_init_func_)
+        if pretrained:
+            self.extract.load_pretrained_()
+        if weights_dict:
+            if len(weights_dict.items()) == len(self.state_dict()):
+                print('Reload all model parameters from weights dict')
+                self.load_state_dict(weights_dict)
+            else:
+                print('Reload part of model parameters from weights dict')
+                self.load_state_dict(weights_dict, strict=False)       
+                
+    def load_batch_(self, batch, dtype='pair'):
+        im_src = batch['src_im'].to(self.device)
+        im_pos = batch['pos_im'].to(self.device)
+        Fs = batch['F'].to(self.device)            
+        if dtype == 'triplet':
+            im_neg = batch['neg_im'].to(self.device)
+            return im_src, im_pos, im_neg, Fs
+        return im_src, im_pos, Fs    
+
+    def forward_coarse_match(self, feat1, feat2, ksize=1):
+        # Feature normalization
+        feat1 = L2Normalize(feat1, dim=1)
+        feat2 = L2Normalize(feat2, dim=1)
+        
+        # Feature correlation
+        corr4d = self.combine(feat1, feat2)
+        
+        # Do 4d maxpooling for relocalization
+        delta4d = None
+        if ksize > 1:
+            corr4d, max_i, max_j, max_k, max_l = maxpool4d(corr4d, k_size=ksize)
+            delta4d = (max_i,max_j,max_k,max_l)
+        corr4d = MutualMatching(corr4d)            
+        corr4d = self.ncn(corr4d)
+        corr4d = MutualMatching(corr4d)        
+        return corr4d, delta4d
+
+    def parse_regressor_out(self, out, psize, ptype, imatches, max_val):
+        w1, h1, w2, h2 = max_val
+        offset = out[:, :4] # N, 4
+        offset = psize * torch.tanh(nn.functional.relu(offset))        
+        if ptype == 'center':
+            shift = psize // 2
+            offset -= shift
+        fmatches = imatches.float() + offset
+        io_probs = out[:, 4]
+        io_probs = torch.sigmoid(io_probs)
+        
+        # Prevent out of range
+        x1 = fmatches[:, 0].clamp(min=0, max=w1)
+        y1 = fmatches[:, 1].clamp(min=0, max=h1)
+        x2 = fmatches[:, 2].clamp(min=0, max=w2)
+        y2 = fmatches[:, 3].clamp(min=0, max=h2)
+        fmatches = torch.stack([x1, y1, x2, y2], dim=-1)
+        return fmatches, io_probs
+
+    def forward_fine_match_mini_batch(self, feats1, feats2, ibatch, imatches,
+                                      psize, ptype, regressor):
+        # ibatch: to index the feature map
+        # imatches: input coarse matches, N, 4
+        N = imatches.shape[0]
+        _, _, h1, w1 = feats1[0].shape
+        _, _, h2, w2 = feats2[0].shape
+        max_val = [w1, h1, w2, h2]
+        
+        f1s, f2s, _, _ = select_local_patch_feats(feats1, feats2,
+                                                  ibatch, imatches,
+                                                  feat_idx=self.feat_idx,
+                                                  feats_downsample=self.feats_downsample,
+                                                  psize=psize,
+                                                  ptype=ptype)    
+        # Feature normalization
+        f1s = L2Normalize(f1s, dim=0)  # D, N*psize*psize
+        f2s = L2Normalize(f2s, dim=0)  # D, N*psize*psize
+
+        # Reshaping: -> (D, N, psize, psize) -> (N, D, psize, psize)
+        f1s = f1s.view(-1, N, psize, psize).permute(1, 0, 2, 3)
+        f2s = f2s.view(-1, N, psize, psize).permute(1, 0, 2, 3)   
+
+     
+        # From im1 to im2       
+        out = regressor(f1s, f2s)  # N, 5
+        fmatches, io_probs = self.parse_regressor_out(out, psize, ptype, imatches, max_val)
+        return fmatches, io_probs
+    
+    def forward_fine_match(self, feats1, feats2, coarse_matches, 
+                           psize, ptype, regressor):
+        batch_size = self.regr_batch
+        masks = []
+        fine_matches = []
+        for ibatch, imatches in enumerate(coarse_matches):
+            # Use mini-batch if too many matches
+            N = imatches.shape[0]
+            if N > batch_size:                
+                batch_inds = [batch_size*i for i in range(N // batch_size + 1)]
+                if batch_inds[-1] < N:
+                    if N -  batch_inds[-1] == 1:
+                        # Special case, slicing leads to 1-dim missing
+                        batch_inds[-1] = N
+                    else:
+                        batch_inds += [N]
+                fmatches = []
+                io_probs = []
+                for bi, (ist, ied) in enumerate(zip(batch_inds[0:-1], batch_inds[1::])):
+                    mini_results = self.forward_fine_match_mini_batch(feats1, feats2, 
+                                                                      ibatch, imatches[ist:ied],
+                                                                      psize, ptype, regressor)
+                    fmatches.append(mini_results[0])
+                    io_probs.append(mini_results[1])
+                fmatches = torch.cat(fmatches, dim=0).squeeze()
+                io_probs = torch.cat(io_probs, dim=0).squeeze()
+            else:
+                fmatches, io_probs = self.forward_fine_match_mini_batch(feats1, feats2, 
+                                                                        ibatch, imatches,
+                                                                        psize, ptype, regressor)
+            fine_matches.append(fmatches)
+            masks.append(io_probs)                        
+        return fine_matches, masks
+
+    def forward(self, im1, im2, ksize=1, return_feats=False):
+        if return_feats:
+            feat1s=[]
+            feat2s=[]
+            self.extract.forward_all(im1, feat1s, early_feat=True)        
+            self.extract.forward_all(im2, feat2s, early_feat=True)
+            feat1 = feat1s[-1]
+            feat2 = feat2s[-1]
+        else:
+            feat1 = self.extract(im1, early_feat=True)        
+            feat2 = self.extract(im2, early_feat=True)     # Shared weights
+        
+        corr4d, delta4d = self.forward_coarse_match(feat1, feat2, ksize=ksize)
+        
+        if return_feats:
+            return corr4d, delta4d, feat1s, feat2s                            
+        else:
+            return corr4d, delta4d
+        
+        
+    def predict_coarse(self, im1, im2, ksize=2, ncn_thres=0.0,
+                       mutual=False, center=True):
+        corr4d, delta4d = self.forward(im1, im2, ksize)
+        coarse_matches, match_scores = self.cal_coarse_matches(corr4d, delta4d, ksize=ksize, 
+                                                               upsample=self.upsample, center=center)
+        
+        # Filter coarse matches
+        coarse_matches, match_scores = filter_coarse(coarse_matches, match_scores, ncn_thres, mutual)
+        return coarse_matches, match_scores
+    
+    def predict_fine(self, im1, im2, ksize=2, ncn_thres=0.0,
+                     mutual=True, return_all=False):
+        corr4d, delta4d, feats1, feats2 = self.forward(im1, im2, ksize=ksize, return_feats=True)
+        coarse_matches, match_scores = self.cal_coarse_matches(corr4d, delta4d, ksize=ksize,
+                                                               upsample=self.upsample, center=True)
+        # Filter coarse matches
+        coarse_matches, match_scores = filter_coarse(coarse_matches, match_scores, ncn_thres, mutual)
+        
+        # Locate initial anchors
+        coarse_matches = self.shift_to_anchors(coarse_matches)
+        
+        # Mid level matching
+        mid_matches, mid_scores = self.forward_fine_match(feats1, feats2, 
+                                                          coarse_matches, 
+                                                          psize=self.psize[0],
+                                                          ptype=self.ptype[0],
+                                                          regressor=self.regress_mid)
+        
+        # Fine level matching
+        fine_matches, fine_scores = self.forward_fine_match(feats1, feats2, 
+                                                            mid_matches,
+                                                            psize=self.psize[1],
+                                                            ptype=self.ptype[1],
+                                                            regressor=self.regress_fine)
+        if return_all:
+            return fine_matches, fine_scores, mid_matches, mid_scores, coarse_matches 
+        return fine_matches, fine_scores, coarse_matches  
+    
+    def refine_matches(self, im1, im2, coarse_matches, io_thres):
+        # Handle empty coarse matches
+        if len(coarse_matches) == 0:
+            return np.empty((0, 4)), np.empty((0,)), np.empty((0, 4))
+
+        if type(coarse_matches) == np.ndarray:
+            coarse_matches_ = torch.from_numpy(coarse_matches).to(self.device).unsqueeze(0)  # 1, N, 4            
+        elif type(coarse_matches) == torch.Tensor:
+            coarse_matches_ = coarse_matches.unsqueeze(0)  # 1, N, 4
+            coarse_matches = coarse_matches.cpu().data.numpy()
+        
+        # Extract local features
+        feat1s=[]
+        feat2s=[]
+        self.extract.forward_all(im1, feat1s, early_feat=True)        
+        self.extract.forward_all(im2, feat2s, early_feat=True)    
+
+        # Mid level matching
+        mid_matches, mid_scores = self.forward_fine_match(feat1s, feat2s, 
+                                                          coarse_matches_, 
+                                                          psize=self.psize[0],
+                                                          ptype=self.ptype[0],
+                                                          regressor=self.regress_mid)
+
+        # Fine level matching
+        fine_matches, fine_scores = self.forward_fine_match(feat1s, feat2s, 
+                                                            mid_matches,
+                                                            psize=self.psize[1],
+                                                            ptype=self.ptype[1],
+                                                            regressor=self.regress_fine)
+        refined_matches = fine_matches[0].cpu().data.numpy()
+        scores = fine_scores[0].cpu().data.numpy()
+        
+        # Further filtering with threshold
+        if io_thres > 0:
+            pos_ids = np.where(scores > io_thres)[0]
+            if len(pos_ids) > 0:
+                coarse_matches = coarse_matches[pos_ids]
+                refined_matches = refined_matches[pos_ids]
+                scores = scores[pos_ids]
+        return refined_matches, scores, coarse_matches
+    
+    def cal_coarse_score(self, corr4d, normalize='softmax'):
+        if normalize is None:
+            normalize = lambda x: x
+        elif normalize == 'softmax':     
+            normalize = lambda x: nn.functional.softmax(x, 1)
+        elif normalize == 'l1':
+            normalize = lambda x: x / (torch.sum(x, dim=1, keepdim=True) + 0.0001)
+        
+        # Mutual matching score
+        batch_size, _, h1, w1, h2, w2 = corr4d.shape
+        nc_B_Avec=corr4d.view(batch_size, h1*w1, h2, w2)
+        nc_A_Bvec=corr4d.view(batch_size, h1, w1, h2*w2).permute(0,3,1,2) # 
+        nc_B_Avec = normalize(nc_B_Avec)
+        nc_A_Bvec = normalize(nc_A_Bvec)
+        scores_B,_= torch.max(nc_B_Avec, dim=1)
+        scores_A,_= torch.max(nc_A_Bvec, dim=1)
+        scores_AB = torch.cat([scores_A.view(-1, h1*w1), scores_B.view(-1, h2*w2)], dim=1)
+        score = scores_AB.mean()
+        return score
+    
+    def cal_coarse_matches(self, corr4d, delta4d, ksize=1, do_softmax=True,
+                           upsample=16, sort=False, center=True, pshift=0):
+        
+        # Original nc implementation: only max locations
+        (xA_, yA_, xB_, yB_, score_) = corr_to_matches(corr4d, delta4d=delta4d,
+                                                       do_softmax=do_softmax,
+                                                       ksize=ksize)
+        (xA2_, yA2_, xB2_, yB2_, score2_) = corr_to_matches(corr4d, delta4d=delta4d,
+                                                            do_softmax=do_softmax,
+                                                            ksize=ksize,
+                                                            invert_matching_direction=True)
+        xA_ = torch.cat((xA_, xA2_), 1)
+        yA_ = torch.cat((yA_, yA2_), 1)
+        xB_ = torch.cat((xB_, xB2_), 1)
+        yB_ = torch.cat((yB_, yB2_), 1)
+        score_ = torch.cat((score_, score2_),1)
+        
+        # Sort as descend
+        if sort:
+            sorted_index = torch.sort(-score_)[1]
+            xA_ = torch.gather(xA_, 1, sorted_index)  # B, 1, N
+            yA_ = torch.gather(yA_, 1, sorted_index)
+            xB_ = torch.gather(xB_, 1, sorted_index)
+            yB_ = torch.gather(yB_, 1, sorted_index)
+            score_ = torch.gather(score_, 1, sorted_index)  # B, N
+
+        xA_ = xA_.unsqueeze(1)
+        yA_ = yA_.unsqueeze(1)
+        xB_ = xB_.unsqueeze(1)
+        yB_ = yB_.unsqueeze(1)        
+        # Create matches and upscale to input resolution
+        matches_ = upsample * torch.cat([xA_, yA_, xB_, yB_], dim=1).permute(0, 2, 1) # B, N, 4
+        if center:
+            delta = upsample // 2
+            matches_ += torch.tensor([[delta, delta, delta, delta]]).unsqueeze(0).to(matches_)                        
+        return matches_, score_
+    
+    def shift_to_anchors(self, matches): 
+        pshift = self.pshift
+        panc = self.panc
+        if panc == 1:
+            return matches
+        
+        # Move pt1/pt2 to its upper-left, upper-right, down-left, down-right
+        # location by pshift, leading to 4 corner anchors
+        # Then take center vs corner from two directions as new matches 
+        shift_template = torch.tensor([
+            [-pshift, -pshift, 0, 0],
+            [pshift, -pshift,  0, 0],
+            [-pshift, pshift,  0, 0],
+            [pshift, pshift,   0, 0],
+            [0, 0, -pshift, -pshift],                
+            [0, 0, pshift, -pshift],
+            [0, 0, -pshift, pshift],
+            [0, 0, pshift, pshift]
+        ]).to(self.device)
+
+        matches_ = []
+        for imatches in matches:
+            imatches =  imatches.unsqueeze(1) + shift_template # N, 16, 4
+            imatches = imatches.reshape(-1, 4)
+            matches_.append(imatches)
+        return matches_
+    

vismatch/third_party/imatch-toolbox/third_party/patch2pix/networks/resnet.py

@@ -0,0 +1,191 @@
+"""
+This script is an adapted version of 
+https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py 
+The goal is to keep ResNet* as only feature extractor, 
+so the code can be used independent of the types of specific tasks,
+i.e., classification or regression. 
+"""
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+from collections import OrderedDict
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = 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:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                               padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = 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:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+    
+class ResNet(nn.Module):
+    PRETRAINED_URLs = {
+        'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+        'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+        'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+        'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+        'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+    }
+    
+    def __init__(self):
+        super().__init__()
+        
+    def _build_model(self, block, layers):
+        self.inplanes = 64
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        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)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+        return nn.Sequential(*layers)
+
+    def forward(self, x, early_feat=False):
+        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)
+        if early_feat:
+            return x
+        x = self.layer4(x)
+        return x
+    
+    def forward_all(self, x, feat_list=[], early_feat=True):
+        feat_list.append(x)
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        feat_list.append(x)
+        
+        x = self.maxpool(x)
+        x = self.layer1(x)
+        feat_list.append(x)
+        
+        x = self.layer2(x)
+        feat_list.append(x)
+        
+        x = self.layer3(x)
+        feat_list.append(x)
+        
+        if not early_feat:
+            x = self.layer4(x)
+            feat_list.append(x)
+    
+    def load_pretrained_(self, ignore='fc'):
+        print('Initialize ResNet using pretrained model from {}'.format(self.pretrained_url))
+        state_dict = model_zoo.load_url(self.pretrained_url)
+        new_state_dict = OrderedDict()
+        for k, v in state_dict.items():
+            if ignore in k:
+                continue
+            new_state_dict[k] = v
+        self.load_state_dict(new_state_dict)
+
+    def change_stride(self, target='layer3'):
+        layer = getattr(self, target)
+        layer[0].conv1.stride = (1, 1)
+        layer[0].conv2.stride = (1, 1)
+        layer[0].downsample[0].stride = (1, 1) 
+
+class ResNet34(ResNet):
+    def __init__(self):
+        super().__init__()
+        self.pretrained_url = self.PRETRAINED_URLs['resnet34']
+        self._build_model(BasicBlock, [3, 4, 6, 3])
+
+class ResNet50(ResNet):
+    def __init__(self):
+        super().__init__()
+        self.pretrained_url = self.PRETRAINED_URLs['resnet50']
+        self._build_model(Bottleneck, [3, 4, 6, 3])
+        
+class ResNet101(ResNet):
+    def __init__(self):
+        super().__init__()
+        self.pretrained_url = self.PRETRAINED_URLs['resnet101']
+        self._build_model(Bottleneck, [3, 4, 23, 3])