what.models.detection.ssd.mobilenet_v2_ssd_lite API documentation

# what.models.detection.ssd.mobilenet_v2_ssd_lite

View Source

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

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

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

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

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

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

    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