what.models.detection.ssd.mobilenet_v1_ssd

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