what.models.detection.ssd.mobilenet_v2_ssd_lite API documentation

# what.models.detection.ssd.mobilenet_v2_ssd_lite

View Source

      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: View Source

     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) View Source

    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): View Source

    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): View Source

    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/'): View Source

     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): View Source

    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