what.models.detection.ssd.mobilenet_v2_ssd_lite

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