what.models.detection.ssd.ssd.ssd

  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):
 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)
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]:
 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):
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):
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):
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):
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):
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):
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:
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)
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