what.models.detection.ssd.mobilenet_v1_ssd

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