what.models.detection.ssd.mobilenet_v2_ssd_lite 

  1import os
  2import itertools
  3
  4import torch
  5from torch.optim.lr_scheduler import CosineAnnealingLR
  6
  7import what.utils.logger as log
  8
  9from .ssd import mobilenet_ssd_config
 10from .ssd.multibox_loss import MultiboxLoss
 11from .ssd.mobilenet_v2_ssd_lite_create import create_mobilenet_v2_ssd_lite, create_mobilenet_v2_ssd_lite_predictor
 12
 13from .utils.misc import freeze_net_layers
 14
 15logger = log.get_logger(__name__)
 16
 17from what.models.detection.datasets.voc import VOC_CLASS_NAMES
 18
 19class MobileNetV2SSDLite:
 20    def __init__(self, model_path, class_names, width_mult = 1.0, is_test=False, device=None):
 21
 22        self.class_names = class_names
 23
 24        self.net = create_mobilenet_v2_ssd_lite(len(self.class_names), is_test=is_test, width_mult=width_mult)
 25
 26        self.predictor = None;
 27        self.device = device;
 28
 29        self.net.load(model_path)
 30
 31    def predict(self, image, top_k=-1, prob_threshold=None):
 32        if self.predictor is None:
 33            self.predictor = create_mobilenet_v2_ssd_lite_predictor(self.net, candidate_size=200, device=self.device)
 34
 35        return self.predictor.predict(image, top_k, prob_threshold)
 36
 37    def step(self, loader, criterion, optimizer, device, debug_steps=100, epoch=-1):
 38        self.net.train(True)
 39        running_loss = 0.0
 40        running_regression_loss = 0.0
 41        running_classification_loss = 0.0
 42        for i, data in enumerate(loader):
 43            images, boxes, labels = data
 44            images = images.to(device)
 45            boxes = boxes.to(device)
 46            labels = labels.to(device)
 47
 48            optimizer.zero_grad()
 49            confidence, locations = self.net(images)
 50            regression_loss, classification_loss = criterion(confidence, locations, labels, boxes)  # TODO CHANGE BOXES
 51            loss = regression_loss + classification_loss
 52            loss.backward()
 53            optimizer.step()
 54
 55            running_loss += loss.item()
 56            running_regression_loss += regression_loss.item()
 57            running_classification_loss += classification_loss.item()
 58            if i and i % debug_steps == 0:
 59                avg_loss = running_loss / debug_steps
 60                avg_reg_loss = running_regression_loss / debug_steps
 61                avg_clf_loss = running_classification_loss / debug_steps
 62                logger.info(
 63                    f"Epoch: {epoch}, Step: {i}, " +
 64                    f"Average Loss: {avg_loss:.4f}, " +
 65                    f"Average Regression Loss {avg_reg_loss:.4f}, " +
 66                    f"Average Classification Loss: {avg_clf_loss:.4f}"
 67                )
 68                running_loss = 0.0
 69                running_regression_loss = 0.0
 70                running_classification_loss = 0.0
 71
 72    def train(self, train_loader, val_loader, device = "cpu", 
 73            scheduler = None, criterion = None, optimizer = None, 
 74            lr = 1e-3, base_net_lr = 1e-3, extra_layers_lr = 1e-3, num_epochs = 100, momentum = 0.9, weight_decay = 5e-4, 
 75            debug_steps = 100, validation_epochs = 5,
 76            freeze_base_net = False, freeze_net = False,
 77            resume = None, base_net = None, pretrained_ssd = None,
 78            checkpoint_folder = "models/"):
 79
 80        if freeze_base_net:
 81            logger.info("Freeze base net.")
 82
 83            freeze_net_layers(self.net.base_net)
 84            params = itertools.chain(self.net.source_layer_add_ons.parameters(), self.net.extras.parameters(),
 85                                    self.net.regression_headers.parameters(), self.net.classification_headers.parameters())
 86            params = [
 87                {'params': itertools.chain(
 88                    self.net.source_layer_add_ons.parameters(),
 89                    self.net.extras.parameters()
 90                ), 'lr': extra_layers_lr},
 91                {'params': itertools.chain(
 92                    self.net.regression_headers.parameters(),
 93                    self.net.classification_headers.parameters()
 94                )}
 95            ]
 96        elif freeze_net:
 97            freeze_net_layers(self.net.base_net)
 98            freeze_net_layers(self.net.source_layer_add_ons)
 99            freeze_net_layers(self.net.extras)
100            params = itertools.chain(self.net.regression_headers.parameters(), self.net.classification_headers.parameters())
101            logger.info("Freeze all the layers except prediction heads.")
102        else:
103            params = [
104                {'params': self.net.base_net.parameters(), 'lr': base_net_lr},
105                {'params': itertools.chain(
106                    self.net.source_layer_add_ons.parameters(),
107                    self.net.extras.parameters()
108                ), 'lr': extra_layers_lr},
109                {'params': itertools.chain(
110                    self.net.regression_headers.parameters(),
111                    self.net.classification_headers.parameters()
112                )}
113            ]
114
115        if resume:
116            logger.info(f"Resume from the model {resume}")
117            self.net.load(resume)
118        elif base_net:
119            logger.info(f"Init from base net {base_net}")
120            self.net.init_from_base_net(base_net)
121        elif pretrained_ssd:
122            logger.info(f"Init from pretrained ssd {pretrained_ssd}")
123            self.net.init_from_pretrained_ssd(pretrained_ssd)
124
125        self.net.to(device)
126
127        if criterion is None:
128            criterion = MultiboxLoss(mobilenet_ssd_config.priors, iou_threshold=0.5, neg_pos_ratio=3,
129                                    center_variance=0.1, size_variance=0.2, device=device)
130        if optimizer is None:
131            optimizer = torch.optim.SGD(params, lr=lr, momentum=momentum,
132                                        weight_decay=weight_decay)
133        if scheduler is None:
134            scheduler = CosineAnnealingLR(optimizer, 120)
135
136        logger.info("Start training using CosineAnnealingLR scheduler.")
137
138        for epoch in range(0, num_epochs):
139            self.step(train_loader, criterion, optimizer, epoch=epoch,
140                device=device, debug_steps=debug_steps)
141
142            scheduler.step()
143
144            if (epoch % validation_epochs == 0) or (epoch == num_epochs - 1):
145                val_loss, val_regression_loss, val_classification_loss = self.eval(val_loader, criterion, device)
146                logger.info(
147                    f"Epoch: {epoch}, " +
148                    f"Validation Loss: {val_loss:.4f}, " +
149                    f"Validation Regression Loss {val_regression_loss:.4f}, " +
150                    f"Validation Classification Loss: {val_classification_loss:.4f}"
151                )
152                model_path = os.path.join(checkpoint_folder, f"mobilenet-v2-ssd-lite-Epoch-{epoch}-Loss-{val_loss}.pth")
153                self.net.save(model_path)
154
155                logger.info(f"Saved model {model_path}")
156
157    def eval(self, loader, criterion, device):
158        self.net.eval()
159        running_loss = 0.0
160        running_regression_loss = 0.0
161        running_classification_loss = 0.0
162        num = 0
163        for _, data in enumerate(loader):
164            images, boxes, labels = data
165            images = images.to(device)
166            boxes = boxes.to(device)
167            labels = labels.to(device)
168            num += 1
169
170            with torch.no_grad():
171                confidence, locations = self.net(images)
172                regression_loss, classification_loss = criterion(confidence, locations, labels, boxes)
173                loss = regression_loss + classification_loss
174
175            running_loss += loss.item()
176            running_regression_loss += regression_loss.item()
177            running_classification_loss += classification_loss.item()
178        return running_loss / num, running_regression_loss / num, running_classification_loss / num
class MobileNetV2SSDLite: 
 20class MobileNetV2SSDLite:
 21    def __init__(self, model_path, class_names, width_mult = 1.0, is_test=False, device=None):
 22
 23        self.class_names = class_names
 24
 25        self.net = create_mobilenet_v2_ssd_lite(len(self.class_names), is_test=is_test, width_mult=width_mult)
 26
 27        self.predictor = None;
 28        self.device = device;
 29
 30        self.net.load(model_path)
 31
 32    def predict(self, image, top_k=-1, prob_threshold=None):
 33        if self.predictor is None:
 34            self.predictor = create_mobilenet_v2_ssd_lite_predictor(self.net, candidate_size=200, device=self.device)
 35
 36        return self.predictor.predict(image, top_k, prob_threshold)
 37
 38    def step(self, loader, criterion, optimizer, device, debug_steps=100, epoch=-1):
 39        self.net.train(True)
 40        running_loss = 0.0
 41        running_regression_loss = 0.0
 42        running_classification_loss = 0.0
 43        for i, data in enumerate(loader):
 44            images, boxes, labels = data
 45            images = images.to(device)
 46            boxes = boxes.to(device)
 47            labels = labels.to(device)
 48
 49            optimizer.zero_grad()
 50            confidence, locations = self.net(images)
 51            regression_loss, classification_loss = criterion(confidence, locations, labels, boxes)  # TODO CHANGE BOXES
 52            loss = regression_loss + classification_loss
 53            loss.backward()
 54            optimizer.step()
 55
 56            running_loss += loss.item()
 57            running_regression_loss += regression_loss.item()
 58            running_classification_loss += classification_loss.item()
 59            if i and i % debug_steps == 0:
 60                avg_loss = running_loss / debug_steps
 61                avg_reg_loss = running_regression_loss / debug_steps
 62                avg_clf_loss = running_classification_loss / debug_steps
 63                logger.info(
 64                    f"Epoch: {epoch}, Step: {i}, " +
 65                    f"Average Loss: {avg_loss:.4f}, " +
 66                    f"Average Regression Loss {avg_reg_loss:.4f}, " +
 67                    f"Average Classification Loss: {avg_clf_loss:.4f}"
 68                )
 69                running_loss = 0.0
 70                running_regression_loss = 0.0
 71                running_classification_loss = 0.0
 72
 73    def train(self, train_loader, val_loader, device = "cpu", 
 74            scheduler = None, criterion = None, optimizer = None, 
 75            lr = 1e-3, base_net_lr = 1e-3, extra_layers_lr = 1e-3, num_epochs = 100, momentum = 0.9, weight_decay = 5e-4, 
 76            debug_steps = 100, validation_epochs = 5,
 77            freeze_base_net = False, freeze_net = False,
 78            resume = None, base_net = None, pretrained_ssd = None,
 79            checkpoint_folder = "models/"):
 80
 81        if freeze_base_net:
 82            logger.info("Freeze base net.")
 83
 84            freeze_net_layers(self.net.base_net)
 85            params = itertools.chain(self.net.source_layer_add_ons.parameters(), self.net.extras.parameters(),
 86                                    self.net.regression_headers.parameters(), self.net.classification_headers.parameters())
 87            params = [
 88                {'params': itertools.chain(
 89                    self.net.source_layer_add_ons.parameters(),
 90                    self.net.extras.parameters()
 91                ), 'lr': extra_layers_lr},
 92                {'params': itertools.chain(
 93                    self.net.regression_headers.parameters(),
 94                    self.net.classification_headers.parameters()
 95                )}
 96            ]
 97        elif freeze_net:
 98            freeze_net_layers(self.net.base_net)
 99            freeze_net_layers(self.net.source_layer_add_ons)
100            freeze_net_layers(self.net.extras)
101            params = itertools.chain(self.net.regression_headers.parameters(), self.net.classification_headers.parameters())
102            logger.info("Freeze all the layers except prediction heads.")
103        else:
104            params = [
105                {'params': self.net.base_net.parameters(), 'lr': base_net_lr},
106                {'params': itertools.chain(
107                    self.net.source_layer_add_ons.parameters(),
108                    self.net.extras.parameters()
109                ), 'lr': extra_layers_lr},
110                {'params': itertools.chain(
111                    self.net.regression_headers.parameters(),
112                    self.net.classification_headers.parameters()
113                )}
114            ]
115
116        if resume:
117            logger.info(f"Resume from the model {resume}")
118            self.net.load(resume)
119        elif base_net:
120            logger.info(f"Init from base net {base_net}")
121            self.net.init_from_base_net(base_net)
122        elif pretrained_ssd:
123            logger.info(f"Init from pretrained ssd {pretrained_ssd}")
124            self.net.init_from_pretrained_ssd(pretrained_ssd)
125
126        self.net.to(device)
127
128        if criterion is None:
129            criterion = MultiboxLoss(mobilenet_ssd_config.priors, iou_threshold=0.5, neg_pos_ratio=3,
130                                    center_variance=0.1, size_variance=0.2, device=device)
131        if optimizer is None:
132            optimizer = torch.optim.SGD(params, lr=lr, momentum=momentum,
133                                        weight_decay=weight_decay)
134        if scheduler is None:
135            scheduler = CosineAnnealingLR(optimizer, 120)
136
137        logger.info("Start training using CosineAnnealingLR scheduler.")
138
139        for epoch in range(0, num_epochs):
140            self.step(train_loader, criterion, optimizer, epoch=epoch,
141                device=device, debug_steps=debug_steps)
142
143            scheduler.step()
144
145            if (epoch % validation_epochs == 0) or (epoch == num_epochs - 1):
146                val_loss, val_regression_loss, val_classification_loss = self.eval(val_loader, criterion, device)
147                logger.info(
148                    f"Epoch: {epoch}, " +
149                    f"Validation Loss: {val_loss:.4f}, " +
150                    f"Validation Regression Loss {val_regression_loss:.4f}, " +
151                    f"Validation Classification Loss: {val_classification_loss:.4f}"
152                )
153                model_path = os.path.join(checkpoint_folder, f"mobilenet-v2-ssd-lite-Epoch-{epoch}-Loss-{val_loss}.pth")
154                self.net.save(model_path)
155
156                logger.info(f"Saved model {model_path}")
157
158    def eval(self, loader, criterion, device):
159        self.net.eval()
160        running_loss = 0.0
161        running_regression_loss = 0.0
162        running_classification_loss = 0.0
163        num = 0
164        for _, data in enumerate(loader):
165            images, boxes, labels = data
166            images = images.to(device)
167            boxes = boxes.to(device)
168            labels = labels.to(device)
169            num += 1
170
171            with torch.no_grad():
172                confidence, locations = self.net(images)
173                regression_loss, classification_loss = criterion(confidence, locations, labels, boxes)
174                loss = regression_loss + classification_loss
175
176            running_loss += loss.item()
177            running_regression_loss += regression_loss.item()
178            running_classification_loss += classification_loss.item()
179        return running_loss / num, running_regression_loss / num, running_classification_loss / num
MobileNetV2SSDLite(model_path, class_names, width_mult=1.0, is_test=False, device=None) 
21    def __init__(self, model_path, class_names, width_mult = 1.0, is_test=False, device=None):
22
23        self.class_names = class_names
24
25        self.net = create_mobilenet_v2_ssd_lite(len(self.class_names), is_test=is_test, width_mult=width_mult)
26
27        self.predictor = None;
28        self.device = device;
29
30        self.net.load(model_path)
def predict(self, image, top_k=-1, prob_threshold=None): 
32    def predict(self, image, top_k=-1, prob_threshold=None):
33        if self.predictor is None:
34            self.predictor = create_mobilenet_v2_ssd_lite_predictor(self.net, candidate_size=200, device=self.device)
35
36        return self.predictor.predict(image, top_k, prob_threshold)
def step( self, loader, criterion, optimizer, device, debug_steps=100, epoch=-1): 
38    def step(self, loader, criterion, optimizer, device, debug_steps=100, epoch=-1):
39        self.net.train(True)
40        running_loss = 0.0
41        running_regression_loss = 0.0
42        running_classification_loss = 0.0
43        for i, data in enumerate(loader):
44            images, boxes, labels = data
45            images = images.to(device)
46            boxes = boxes.to(device)
47            labels = labels.to(device)
48
49            optimizer.zero_grad()
50            confidence, locations = self.net(images)
51            regression_loss, classification_loss = criterion(confidence, locations, labels, boxes)  # TODO CHANGE BOXES
52            loss = regression_loss + classification_loss
53            loss.backward()
54            optimizer.step()
55
56            running_loss += loss.item()
57            running_regression_loss += regression_loss.item()
58            running_classification_loss += classification_loss.item()
59            if i and i % debug_steps == 0:
60                avg_loss = running_loss / debug_steps
61                avg_reg_loss = running_regression_loss / debug_steps
62                avg_clf_loss = running_classification_loss / debug_steps
63                logger.info(
64                    f"Epoch: {epoch}, Step: {i}, " +
65                    f"Average Loss: {avg_loss:.4f}, " +
66                    f"Average Regression Loss {avg_reg_loss:.4f}, " +
67                    f"Average Classification Loss: {avg_clf_loss:.4f}"
68                )
69                running_loss = 0.0
70                running_regression_loss = 0.0
71                running_classification_loss = 0.0
def train( self, train_loader, val_loader, device='cpu', scheduler=None, criterion=None, optimizer=None, lr=0.001, base_net_lr=0.001, extra_layers_lr=0.001, num_epochs=100, momentum=0.9, weight_decay=0.0005, debug_steps=100, validation_epochs=5, freeze_base_net=False, freeze_net=False, resume=None, base_net=None, pretrained_ssd=None, checkpoint_folder='models/'): 
 73    def train(self, train_loader, val_loader, device = "cpu", 
 74            scheduler = None, criterion = None, optimizer = None, 
 75            lr = 1e-3, base_net_lr = 1e-3, extra_layers_lr = 1e-3, num_epochs = 100, momentum = 0.9, weight_decay = 5e-4, 
 76            debug_steps = 100, validation_epochs = 5,
 77            freeze_base_net = False, freeze_net = False,
 78            resume = None, base_net = None, pretrained_ssd = None,
 79            checkpoint_folder = "models/"):
 80
 81        if freeze_base_net:
 82            logger.info("Freeze base net.")
 83
 84            freeze_net_layers(self.net.base_net)
 85            params = itertools.chain(self.net.source_layer_add_ons.parameters(), self.net.extras.parameters(),
 86                                    self.net.regression_headers.parameters(), self.net.classification_headers.parameters())
 87            params = [
 88                {'params': itertools.chain(
 89                    self.net.source_layer_add_ons.parameters(),
 90                    self.net.extras.parameters()
 91                ), 'lr': extra_layers_lr},
 92                {'params': itertools.chain(
 93                    self.net.regression_headers.parameters(),
 94                    self.net.classification_headers.parameters()
 95                )}
 96            ]
 97        elif freeze_net:
 98            freeze_net_layers(self.net.base_net)
 99            freeze_net_layers(self.net.source_layer_add_ons)
100            freeze_net_layers(self.net.extras)
101            params = itertools.chain(self.net.regression_headers.parameters(), self.net.classification_headers.parameters())
102            logger.info("Freeze all the layers except prediction heads.")
103        else:
104            params = [
105                {'params': self.net.base_net.parameters(), 'lr': base_net_lr},
106                {'params': itertools.chain(
107                    self.net.source_layer_add_ons.parameters(),
108                    self.net.extras.parameters()
109                ), 'lr': extra_layers_lr},
110                {'params': itertools.chain(
111                    self.net.regression_headers.parameters(),
112                    self.net.classification_headers.parameters()
113                )}
114            ]
115
116        if resume:
117            logger.info(f"Resume from the model {resume}")
118            self.net.load(resume)
119        elif base_net:
120            logger.info(f"Init from base net {base_net}")
121            self.net.init_from_base_net(base_net)
122        elif pretrained_ssd:
123            logger.info(f"Init from pretrained ssd {pretrained_ssd}")
124            self.net.init_from_pretrained_ssd(pretrained_ssd)
125
126        self.net.to(device)
127
128        if criterion is None:
129            criterion = MultiboxLoss(mobilenet_ssd_config.priors, iou_threshold=0.5, neg_pos_ratio=3,
130                                    center_variance=0.1, size_variance=0.2, device=device)
131        if optimizer is None:
132            optimizer = torch.optim.SGD(params, lr=lr, momentum=momentum,
133                                        weight_decay=weight_decay)
134        if scheduler is None:
135            scheduler = CosineAnnealingLR(optimizer, 120)
136
137        logger.info("Start training using CosineAnnealingLR scheduler.")
138
139        for epoch in range(0, num_epochs):
140            self.step(train_loader, criterion, optimizer, epoch=epoch,
141                device=device, debug_steps=debug_steps)
142
143            scheduler.step()
144
145            if (epoch % validation_epochs == 0) or (epoch == num_epochs - 1):
146                val_loss, val_regression_loss, val_classification_loss = self.eval(val_loader, criterion, device)
147                logger.info(
148                    f"Epoch: {epoch}, " +
149                    f"Validation Loss: {val_loss:.4f}, " +
150                    f"Validation Regression Loss {val_regression_loss:.4f}, " +
151                    f"Validation Classification Loss: {val_classification_loss:.4f}"
152                )
153                model_path = os.path.join(checkpoint_folder, f"mobilenet-v2-ssd-lite-Epoch-{epoch}-Loss-{val_loss}.pth")
154                self.net.save(model_path)
155
156                logger.info(f"Saved model {model_path}")
def eval(self, loader, criterion, device): 
158    def eval(self, loader, criterion, device):
159        self.net.eval()
160        running_loss = 0.0
161        running_regression_loss = 0.0
162        running_classification_loss = 0.0
163        num = 0
164        for _, data in enumerate(loader):
165            images, boxes, labels = data
166            images = images.to(device)
167            boxes = boxes.to(device)
168            labels = labels.to(device)
169            num += 1
170
171            with torch.no_grad():
172                confidence, locations = self.net(images)
173                regression_loss, classification_loss = criterion(confidence, locations, labels, boxes)
174                loss = regression_loss + classification_loss
175
176            running_loss += loss.item()
177            running_regression_loss += regression_loss.item()
178            running_classification_loss += classification_loss.item()
179        return running_loss / num, running_regression_loss / num, running_classification_loss / num