what.models.detection.ssd.ssd.ssd API documentation

# what.models.detection.ssd.ssd.ssd

View Source

      1import numpy as np
      2from typing import List, Tuple
      3from collections import namedtuple
      4
      5import torch
      6import torch.nn as nn
      7import torch.nn.functional as F
      8from torch.utils.model_zoo import load_url as load_state_dict_from_url
      9
     10from ..utils import box_utils
     11
     12GraphPath = namedtuple("GraphPath", ['s0', 'name', 's1'])  #
     13
     14class SSD(nn.Module):
     15    def __init__(self, num_classes: int, base_net: nn.ModuleList, source_layer_indexes: List[int],
     16                 extras: nn.ModuleList, classification_headers: nn.ModuleList,
     17                 regression_headers: nn.ModuleList, is_test=False, config=None, device=None):
     18        """Compose a SSD model using the given components.
     19        """
     20        super(SSD, self).__init__()
     21
     22        self.num_classes = num_classes
     23        self.base_net = base_net
     24        self.source_layer_indexes = source_layer_indexes
     25        self.extras = extras
     26        self.classification_headers = classification_headers
     27        self.regression_headers = regression_headers
     28        self.is_test = is_test
     29        self.config = config
     30
     31        # register layers in source_layer_indexes by adding them to a module list
     32        self.source_layer_add_ons = nn.ModuleList([t[1] for t in source_layer_indexes
     33                                                   if isinstance(t, tuple) and not isinstance(t, GraphPath)])
     34        if device:
     35            self.device = device
     36        else:
     37            self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     38        if is_test:
     39            self.config = config
     40            self.priors = config.priors.to(self.device)
     41
     42    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
     43        confidences = []
     44        locations = []
     45        start_layer_index = 0
     46        header_index = 0
     47        for end_layer_index in self.source_layer_indexes:
     48            if isinstance(end_layer_index, GraphPath):
     49                path = end_layer_index
     50                end_layer_index = end_layer_index.s0
     51                added_layer = None
     52            elif isinstance(end_layer_index, tuple):
     53                added_layer = end_layer_index[1]
     54                end_layer_index = end_layer_index[0]
     55                path = None
     56            else:
     57                added_layer = None
     58                path = None
     59            for layer in self.base_net[start_layer_index: end_layer_index]:
     60                x = layer(x)
     61            if added_layer:
     62                y = added_layer(x)
     63            else:
     64                y = x
     65            if path:
     66                sub = getattr(self.base_net[end_layer_index], path.name)
     67                for layer in sub[:path.s1]:
     68                    x = layer(x)
     69                y = x
     70                for layer in sub[path.s1:]:
     71                    x = layer(x)
     72                end_layer_index += 1
     73            start_layer_index = end_layer_index
     74            confidence, location = self.compute_header(header_index, y)
     75            header_index += 1
     76            confidences.append(confidence)
     77            locations.append(location)
     78
     79        for layer in self.base_net[end_layer_index:]:
     80            x = layer(x)
     81
     82        for layer in self.extras:
     83            x = layer(x)
     84            confidence, location = self.compute_header(header_index, x)
     85            header_index += 1
     86            confidences.append(confidence)
     87            locations.append(location)
     88
     89        confidences = torch.cat(confidences, 1)
     90        locations = torch.cat(locations, 1)
     91
     92        if self.is_test:
     93            confidences = F.softmax(confidences, dim=2)
     94            boxes = box_utils.convert_locations_to_boxes(
     95                locations.to(self.device), self.priors.to(self.device), self.config.center_variance, self.config.size_variance
     96            )
     97            boxes = box_utils.center_form_to_corner_form(boxes)
     98            return confidences, boxes
     99        else:
    100            return confidences, locations
    101
    102    def compute_header(self, i, x):
    103        confidence = self.classification_headers[i](x)
    104        confidence = confidence.permute(0, 2, 3, 1).contiguous()
    105        confidence = confidence.view(confidence.size(0), -1, self.num_classes)
    106
    107        location = self.regression_headers[i](x)
    108        location = location.permute(0, 2, 3, 1).contiguous()
    109        location = location.view(location.size(0), -1, 4)
    110
    111        return confidence, location
    112
    113    def init(self):
    114        self.base_net.apply(_xavier_init_)
    115
    116        self.source_layer_add_ons.apply(_xavier_init_)
    117        self.extras.apply(_xavier_init_)
    118
    119        self.classification_headers.apply(_xavier_init_)
    120        self.regression_headers.apply(_xavier_init_)
    121
    122    def init_from_base_net(self, model):
    123        self.base_net.load_state_dict(torch.load(model, map_location=lambda storage, loc: storage), strict=True)
    124
    125        self.source_layer_add_ons.apply(_xavier_init_)
    126        self.extras.apply(_xavier_init_)
    127
    128        self.classification_headers.apply(_xavier_init_)
    129        self.regression_headers.apply(_xavier_init_)
    130
    131    def init_from_pretrained_ssd(self, model):
    132        state_dict = torch.load(model, map_location=lambda storage, loc: storage)
    133        state_dict = {k: v for k, v in state_dict.items() if not (k.startswith("classification_headers") or k.startswith("regression_headers"))}
    134        model_dict = self.state_dict()
    135        model_dict.update(state_dict)
    136        self.load_state_dict(model_dict)
    137
    138        self.classification_headers.apply(_xavier_init_)
    139        self.regression_headers.apply(_xavier_init_)
    140
    141    def load(self, model, pretrained=False):
    142        if pretrained is True:
    143            state_dict = load_state_dict_from_url(model, progress=True, map_location=lambda storage, loc: storage)
    144            self.load_state_dict(state_dict)
    145        else:
    146            self.load_state_dict(torch.load(model, map_location=lambda storage, loc: storage))
    147
    148    def save(self, model_path):
    149        torch.save(self.state_dict(), model_path)
    150
    151
    152class MatchPrior(object):
    153    def __init__(self, center_form_priors, center_variance, size_variance, iou_threshold):
    154        self.center_form_priors = center_form_priors
    155        self.corner_form_priors = box_utils.center_form_to_corner_form(center_form_priors)
    156        self.center_variance = center_variance
    157        self.size_variance = size_variance
    158        self.iou_threshold = iou_threshold
    159
    160    def __call__(self, gt_boxes, gt_labels):
    161        if type(gt_boxes) is np.ndarray:
    162            gt_boxes = torch.from_numpy(gt_boxes)
    163        if type(gt_labels) is np.ndarray:
    164            gt_labels = torch.from_numpy(gt_labels)
    165        boxes, labels = box_utils.assign_priors(gt_boxes, gt_labels,
    166                                                self.corner_form_priors, self.iou_threshold)
    167        boxes = box_utils.corner_form_to_center_form(boxes)
    168        locations = box_utils.convert_boxes_to_locations(boxes, self.center_form_priors, self.center_variance, self.size_variance)
    169        return locations, labels
    170
    171def _xavier_init_(m: nn.Module):
    172    if isinstance(m, nn.Conv2d):
    173        nn.init.xavier_uniform_(m.weight)

class GraphPath(builtins.tuple):

GraphPath(s0, name, s1)

GraphPath(s0, name, s1)

Create new instance of GraphPath(s0, name, s1)

s0

Alias for field number 0

name

Alias for field number 1

s1

Alias for field number 2

##### Inherited Members

builtins.tuple
  index
  count

class SSD(torch.nn.modules.module.Module): View Source

     15class SSD(nn.Module):
     16    def __init__(self, num_classes: int, base_net: nn.ModuleList, source_layer_indexes: List[int],
     17                 extras: nn.ModuleList, classification_headers: nn.ModuleList,
     18                 regression_headers: nn.ModuleList, is_test=False, config=None, device=None):
     19        """Compose a SSD model using the given components.
     20        """
     21        super(SSD, self).__init__()
     22
     23        self.num_classes = num_classes
     24        self.base_net = base_net
     25        self.source_layer_indexes = source_layer_indexes
     26        self.extras = extras
     27        self.classification_headers = classification_headers
     28        self.regression_headers = regression_headers
     29        self.is_test = is_test
     30        self.config = config
     31
     32        # register layers in source_layer_indexes by adding them to a module list
     33        self.source_layer_add_ons = nn.ModuleList([t[1] for t in source_layer_indexes
     34                                                   if isinstance(t, tuple) and not isinstance(t, GraphPath)])
     35        if device:
     36            self.device = device
     37        else:
     38            self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     39        if is_test:
     40            self.config = config
     41            self.priors = config.priors.to(self.device)
     42
     43    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
     44        confidences = []
     45        locations = []
     46        start_layer_index = 0
     47        header_index = 0
     48        for end_layer_index in self.source_layer_indexes:
     49            if isinstance(end_layer_index, GraphPath):
     50                path = end_layer_index
     51                end_layer_index = end_layer_index.s0
     52                added_layer = None
     53            elif isinstance(end_layer_index, tuple):
     54                added_layer = end_layer_index[1]
     55                end_layer_index = end_layer_index[0]
     56                path = None
     57            else:
     58                added_layer = None
     59                path = None
     60            for layer in self.base_net[start_layer_index: end_layer_index]:
     61                x = layer(x)
     62            if added_layer:
     63                y = added_layer(x)
     64            else:
     65                y = x
     66            if path:
     67                sub = getattr(self.base_net[end_layer_index], path.name)
     68                for layer in sub[:path.s1]:
     69                    x = layer(x)
     70                y = x
     71                for layer in sub[path.s1:]:
     72                    x = layer(x)
     73                end_layer_index += 1
     74            start_layer_index = end_layer_index
     75            confidence, location = self.compute_header(header_index, y)
     76            header_index += 1
     77            confidences.append(confidence)
     78            locations.append(location)
     79
     80        for layer in self.base_net[end_layer_index:]:
     81            x = layer(x)
     82
     83        for layer in self.extras:
     84            x = layer(x)
     85            confidence, location = self.compute_header(header_index, x)
     86            header_index += 1
     87            confidences.append(confidence)
     88            locations.append(location)
     89
     90        confidences = torch.cat(confidences, 1)
     91        locations = torch.cat(locations, 1)
     92
     93        if self.is_test:
     94            confidences = F.softmax(confidences, dim=2)
     95            boxes = box_utils.convert_locations_to_boxes(
     96                locations.to(self.device), self.priors.to(self.device), self.config.center_variance, self.config.size_variance
     97            )
     98            boxes = box_utils.center_form_to_corner_form(boxes)
     99            return confidences, boxes
    100        else:
    101            return confidences, locations
    102
    103    def compute_header(self, i, x):
    104        confidence = self.classification_headers[i](x)
    105        confidence = confidence.permute(0, 2, 3, 1).contiguous()
    106        confidence = confidence.view(confidence.size(0), -1, self.num_classes)
    107
    108        location = self.regression_headers[i](x)
    109        location = location.permute(0, 2, 3, 1).contiguous()
    110        location = location.view(location.size(0), -1, 4)
    111
    112        return confidence, location
    113
    114    def init(self):
    115        self.base_net.apply(_xavier_init_)
    116
    117        self.source_layer_add_ons.apply(_xavier_init_)
    118        self.extras.apply(_xavier_init_)
    119
    120        self.classification_headers.apply(_xavier_init_)
    121        self.regression_headers.apply(_xavier_init_)
    122
    123    def init_from_base_net(self, model):
    124        self.base_net.load_state_dict(torch.load(model, map_location=lambda storage, loc: storage), strict=True)
    125
    126        self.source_layer_add_ons.apply(_xavier_init_)
    127        self.extras.apply(_xavier_init_)
    128
    129        self.classification_headers.apply(_xavier_init_)
    130        self.regression_headers.apply(_xavier_init_)
    131
    132    def init_from_pretrained_ssd(self, model):
    133        state_dict = torch.load(model, map_location=lambda storage, loc: storage)
    134        state_dict = {k: v for k, v in state_dict.items() if not (k.startswith("classification_headers") or k.startswith("regression_headers"))}
    135        model_dict = self.state_dict()
    136        model_dict.update(state_dict)
    137        self.load_state_dict(model_dict)
    138
    139        self.classification_headers.apply(_xavier_init_)
    140        self.regression_headers.apply(_xavier_init_)
    141
    142    def load(self, model, pretrained=False):
    143        if pretrained is True:
    144            state_dict = load_state_dict_from_url(model, progress=True, map_location=lambda storage, loc: storage)
    145            self.load_state_dict(state_dict)
    146        else:
    147            self.load_state_dict(torch.load(model, map_location=lambda storage, loc: storage))
    148
    149    def save(self, model_path):
    150        torch.save(self.state_dict(), model_path)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

    import torch.nn as nn
    import torch.nn.functional as F

    class Model(nn.Module):
        def __init__(self):
            super().__init__()
            self.conv1 = nn.Conv2d(1, 20, 5)
            self.conv2 = nn.Conv2d(20, 20, 5)

        def forward(self, x):
            x = F.relu(self.conv1(x))
            return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call `to()`, etc.

As per the example above, an `__init__()` call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

SSD( num_classes: int, base_net: torch.nn.modules.container.ModuleList, source_layer_indexes: List[int], extras: torch.nn.modules.container.ModuleList, classification_headers: torch.nn.modules.container.ModuleList, regression_headers: torch.nn.modules.container.ModuleList, is_test=False, config=None, device=None) View Source

    16    def __init__(self, num_classes: int, base_net: nn.ModuleList, source_layer_indexes: List[int],
    17                 extras: nn.ModuleList, classification_headers: nn.ModuleList,
    18                 regression_headers: nn.ModuleList, is_test=False, config=None, device=None):
    19        """Compose a SSD model using the given components.
    20        """
    21        super(SSD, self).__init__()
    22
    23        self.num_classes = num_classes
    24        self.base_net = base_net
    25        self.source_layer_indexes = source_layer_indexes
    26        self.extras = extras
    27        self.classification_headers = classification_headers
    28        self.regression_headers = regression_headers
    29        self.is_test = is_test
    30        self.config = config
    31
    32        # register layers in source_layer_indexes by adding them to a module list
    33        self.source_layer_add_ons = nn.ModuleList([t[1] for t in source_layer_indexes
    34                                                   if isinstance(t, tuple) and not isinstance(t, GraphPath)])
    35        if device:
    36            self.device = device
    37        else:
    38            self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    39        if is_test:
    40            self.config = config
    41            self.priors = config.priors.to(self.device)

Compose a SSD model using the given components.

def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: View Source

     43    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
     44        confidences = []
     45        locations = []
     46        start_layer_index = 0
     47        header_index = 0
     48        for end_layer_index in self.source_layer_indexes:
     49            if isinstance(end_layer_index, GraphPath):
     50                path = end_layer_index
     51                end_layer_index = end_layer_index.s0
     52                added_layer = None
     53            elif isinstance(end_layer_index, tuple):
     54                added_layer = end_layer_index[1]
     55                end_layer_index = end_layer_index[0]
     56                path = None
     57            else:
     58                added_layer = None
     59                path = None
     60            for layer in self.base_net[start_layer_index: end_layer_index]:
     61                x = layer(x)
     62            if added_layer:
     63                y = added_layer(x)
     64            else:
     65                y = x
     66            if path:
     67                sub = getattr(self.base_net[end_layer_index], path.name)
     68                for layer in sub[:path.s1]:
     69                    x = layer(x)
     70                y = x
     71                for layer in sub[path.s1:]:
     72                    x = layer(x)
     73                end_layer_index += 1
     74            start_layer_index = end_layer_index
     75            confidence, location = self.compute_header(header_index, y)
     76            header_index += 1
     77            confidences.append(confidence)
     78            locations.append(location)
     79
     80        for layer in self.base_net[end_layer_index:]:
     81            x = layer(x)
     82
     83        for layer in self.extras:
     84            x = layer(x)
     85            confidence, location = self.compute_header(header_index, x)
     86            header_index += 1
     87            confidences.append(confidence)
     88            locations.append(location)
     89
     90        confidences = torch.cat(confidences, 1)
     91        locations = torch.cat(locations, 1)
     92
     93        if self.is_test:
     94            confidences = F.softmax(confidences, dim=2)
     95            boxes = box_utils.convert_locations_to_boxes(
     96                locations.to(self.device), self.priors.to(self.device), self.config.center_variance, self.config.size_variance
     97            )
     98            boxes = box_utils.center_form_to_corner_form(boxes)
     99            return confidences, boxes
    100        else:
    101            return confidences, locations

Defines the computation performed at every call.

Should be overridden by all subclasses.

Although the recipe for forward pass needs to be defined within this function, one should call the `Module` instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

def compute_header(self, i, x): View Source

    103    def compute_header(self, i, x):
    104        confidence = self.classification_headers[i](x)
    105        confidence = confidence.permute(0, 2, 3, 1).contiguous()
    106        confidence = confidence.view(confidence.size(0), -1, self.num_classes)
    107
    108        location = self.regression_headers[i](x)
    109        location = location.permute(0, 2, 3, 1).contiguous()
    110        location = location.view(location.size(0), -1, 4)
    111
    112        return confidence, location

def init(self): View Source

    114    def init(self):
    115        self.base_net.apply(_xavier_init_)
    116
    117        self.source_layer_add_ons.apply(_xavier_init_)
    118        self.extras.apply(_xavier_init_)
    119
    120        self.classification_headers.apply(_xavier_init_)
    121        self.regression_headers.apply(_xavier_init_)

def init_from_base_net(self, model): View Source

    123    def init_from_base_net(self, model):
    124        self.base_net.load_state_dict(torch.load(model, map_location=lambda storage, loc: storage), strict=True)
    125
    126        self.source_layer_add_ons.apply(_xavier_init_)
    127        self.extras.apply(_xavier_init_)
    128
    129        self.classification_headers.apply(_xavier_init_)
    130        self.regression_headers.apply(_xavier_init_)

def init_from_pretrained_ssd(self, model): View Source

    132    def init_from_pretrained_ssd(self, model):
    133        state_dict = torch.load(model, map_location=lambda storage, loc: storage)
    134        state_dict = {k: v for k, v in state_dict.items() if not (k.startswith("classification_headers") or k.startswith("regression_headers"))}
    135        model_dict = self.state_dict()
    136        model_dict.update(state_dict)
    137        self.load_state_dict(model_dict)
    138
    139        self.classification_headers.apply(_xavier_init_)
    140        self.regression_headers.apply(_xavier_init_)

def load(self, model, pretrained=False): View Source

    142    def load(self, model, pretrained=False):
    143        if pretrained is True:
    144            state_dict = load_state_dict_from_url(model, progress=True, map_location=lambda storage, loc: storage)
    145            self.load_state_dict(state_dict)
    146        else:
    147            self.load_state_dict(torch.load(model, map_location=lambda storage, loc: storage))

def save(self, model_path): View Source

    149    def save(self, model_path):
    150        torch.save(self.state_dict(), model_path)

##### Inherited Members

torch.nn.modules.module.Module
  dump_patches
  register_buffer
  register_parameter
  add_module
  register_module
  get_submodule
  get_parameter
  get_buffer
  get_extra_state
  set_extra_state
  apply
  cuda
  xpu
  cpu
  type
  float
  double
  half
  bfloat16
  to_empty
  to
  register_backward_hook
  register_full_backward_hook
  register_forward_pre_hook
  register_forward_hook
  state_dict
  load_state_dict
  parameters
  named_parameters
  buffers
  named_buffers
  children
  named_children
  modules
  named_modules
  train
  eval
  requires_grad_
  zero_grad
  share_memory
  extra_repr

class MatchPrior: View Source

    153class MatchPrior(object):
    154    def __init__(self, center_form_priors, center_variance, size_variance, iou_threshold):
    155        self.center_form_priors = center_form_priors
    156        self.corner_form_priors = box_utils.center_form_to_corner_form(center_form_priors)
    157        self.center_variance = center_variance
    158        self.size_variance = size_variance
    159        self.iou_threshold = iou_threshold
    160
    161    def __call__(self, gt_boxes, gt_labels):
    162        if type(gt_boxes) is np.ndarray:
    163            gt_boxes = torch.from_numpy(gt_boxes)
    164        if type(gt_labels) is np.ndarray:
    165            gt_labels = torch.from_numpy(gt_labels)
    166        boxes, labels = box_utils.assign_priors(gt_boxes, gt_labels,
    167                                                self.corner_form_priors, self.iou_threshold)
    168        boxes = box_utils.corner_form_to_center_form(boxes)
    169        locations = box_utils.convert_boxes_to_locations(boxes, self.center_form_priors, self.center_variance, self.size_variance)
    170        return locations, labels

MatchPrior(center_form_priors, center_variance, size_variance, iou_threshold) View Source

    154    def __init__(self, center_form_priors, center_variance, size_variance, iou_threshold):
    155        self.center_form_priors = center_form_priors
    156        self.corner_form_priors = box_utils.center_form_to_corner_form(center_form_priors)
    157        self.center_variance = center_variance
    158        self.size_variance = size_variance
    159        self.iou_threshold = iou_threshold