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)
GraphPath(s0, name, s1)
Inherited Members
- builtins.tuple
- index
- count
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
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.
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.
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
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_)
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_)
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))
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
- extra_repr
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
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