what.models.detection.ssd.mobilenet_v1_ssd API documentation

# what.models.detection.ssd.mobilenet_v1_ssd

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.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

class MobileNetV1SSD: View Source

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

MobileNetV1SSD(model_path, class_names, is_test=False, device=None) View Source

    19    def __init__(self, model_path, class_names, is_test=False, device=None):
    20
    21        self.class_names = class_names
    22
    23        self.net = create_mobilenet_v1_ssd(len(self.class_names), is_test=is_test)
    24
    25        self.predictor = None;
    26        self.device = device;
    27
    28        self.net.load(model_path)

def predict(self, image, top_k=-1, prob_threshold=None): View Source

    30    def predict(self, image, top_k=-1, prob_threshold=None):
    31        if self.predictor is None:
    32            self.predictor = create_mobilenet_v1_ssd_predictor(self.net, device=self.device, candidate_size=200)
    33
    34        return self.predictor.predict(image, top_k, prob_threshold)

def step( self, loader, criterion, optimizer, device, debug_steps=100, epoch=-1): View Source

    36    def step(self, loader, criterion, optimizer, device, debug_steps=100, epoch=-1):
    37        self.net.train(True)
    38        running_loss = 0.0
    39        running_regression_loss = 0.0
    40        running_classification_loss = 0.0
    41        for i, data in enumerate(loader):
    42            images, boxes, labels = data
    43            images = images.to(device)
    44            boxes = boxes.to(device)
    45            labels = labels.to(device)
    46
    47            optimizer.zero_grad()
    48            confidence, locations = self.net(images)
    49            regression_loss, classification_loss = criterion(confidence, locations, labels, boxes)  # TODO CHANGE BOXES
    50            loss = regression_loss + classification_loss
    51            loss.backward()
    52            optimizer.step()
    53
    54            running_loss += loss.item()
    55            running_regression_loss += regression_loss.item()
    56            running_classification_loss += classification_loss.item()
    57            if i and i % debug_steps == 0:
    58                avg_loss = running_loss / debug_steps
    59                avg_reg_loss = running_regression_loss / debug_steps
    60                avg_clf_loss = running_classification_loss / debug_steps
    61                logger.info(
    62                    f"Epoch: {epoch}, Step: {i}, " +
    63                    f"Average Loss: {avg_loss:.4f}, " +
    64                    f"Average Regression Loss {avg_reg_loss:.4f}, " +
    65                    f"Average Classification Loss: {avg_clf_loss:.4f}"
    66                )
    67                running_loss = 0.0
    68                running_regression_loss = 0.0
    69                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

     71    def train(self, train_loader, val_loader, device = "cpu",
     72             scheduler = None, criterion = None, optimizer = None,
     73             lr = 1e-3, base_net_lr = 1e-3, extra_layers_lr = 1e-3, num_epochs = 100, momentum = 0.9, weight_decay = 5e-4,
     74             debug_steps = 100, validation_epochs = 5,
     75             freeze_base_net = False, freeze_net = False,
     76             resume = None, base_net = None, pretrained_ssd = None,
     77             checkpoint_folder = "models/"):
     78
     79        if freeze_base_net:
     80            logger.info("Freeze base net.")
     81
     82            freeze_net_layers(self.net.base_net)
     83            params = itertools.chain(self.net.source_layer_add_ons.parameters(), self.net.extras.parameters(),
     84                                    self.net.regression_headers.parameters(), self.net.classification_headers.parameters())
     85            params = [
     86                {'params': itertools.chain(
     87                    self.net.source_layer_add_ons.parameters(),
     88                    self.net.extras.parameters()
     89                ), 'lr': extra_layers_lr},
     90                {'params': itertools.chain(
     91                    self.net.regression_headers.parameters(),
     92                    self.net.classification_headers.parameters()
     93                )}
     94            ]
     95        elif freeze_net:
     96            freeze_net_layers(self.net.base_net)
     97            freeze_net_layers(self.net.source_layer_add_ons)
     98            freeze_net_layers(self.net.extras)
     99            params = itertools.chain(self.net.regression_headers.parameters(), self.net.classification_headers.parameters())
    100            logger.info("Freeze all the layers except prediction heads.")
    101        else:
    102            params = [
    103                {'params': self.net.base_net.parameters(), 'lr': base_net_lr},
    104                {'params': itertools.chain(
    105                    self.net.source_layer_add_ons.parameters(),
    106                    self.net.extras.parameters()
    107                ), 'lr': extra_layers_lr},
    108                {'params': itertools.chain(
    109                    self.net.regression_headers.parameters(),
    110                    self.net.classification_headers.parameters()
    111                )}
    112            ]
    113
    114        if resume:
    115            logger.info(f"Resume from the model {resume}")
    116            self.net.load(resume)
    117        elif base_net:
    118            logger.info(f"Init from base net {base_net}")
    119            self.net.init_from_base_net(base_net)
    120        elif pretrained_ssd:
    121            logger.info(f"Init from pretrained ssd {pretrained_ssd}")
    122            self.net.init_from_pretrained_ssd(pretrained_ssd)
    123
    124        self.net.to(device)
    125
    126        if criterion is None:
    127            criterion = MultiboxLoss(mobilenet_ssd_config.priors, iou_threshold=0.5, neg_pos_ratio=3,
    128                                    center_variance=0.1, size_variance=0.2, device=device)
    129        if optimizer is None:
    130            optimizer = torch.optim.SGD(params, lr=lr, momentum=momentum,
    131                                        weight_decay=weight_decay)
    132        if scheduler is None:
    133            scheduler = CosineAnnealingLR(optimizer, 120)
    134
    135        logger.info("Start training using CosineAnnealingLR scheduler.")
    136
    137        for epoch in range(0, num_epochs):
    138            self.step(train_loader, criterion, optimizer, epoch=epoch,
    139                device=device, debug_steps=debug_steps)
    140
    141            scheduler.step()
    142
    143            if (epoch % validation_epochs == 0) or (epoch == num_epochs - 1):
    144                val_loss, val_regression_loss, val_classification_loss = self.eval(val_loader, criterion, device)
    145                logger.info(
    146                    f"Epoch: {epoch}, " +
    147                    f"Validation Loss: {val_loss:.4f}, " +
    148                    f"Validation Regression Loss {val_regression_loss:.4f}, " +
    149                    f"Validation Classification Loss: {val_classification_loss:.4f}"
    150                )
    151                model_path = os.path.join(checkpoint_folder, f"mobilenet-v1-ssd-Epoch-{epoch}-Loss-{val_loss}.pth")
    152                self.net.save(model_path)
    153
    154                logger.info(f"Saved model {model_path}")

def eval(self, loader, criterion, device): View Source

    156    def eval(self, loader, criterion, device):
    157        self.net.eval()
    158        running_loss = 0.0
    159        running_regression_loss = 0.0
    160        running_classification_loss = 0.0
    161        num = 0
    162        for _, data in enumerate(loader):
    163            images, boxes, labels = data
    164            images = images.to(device)
    165            boxes = boxes.to(device)
    166            labels = labels.to(device)
    167            num += 1
    168
    169            with torch.no_grad():
    170                confidence, locations = self.net(images)
    171                regression_loss, classification_loss = criterion(confidence, locations, labels, boxes)
    172                loss = regression_loss + classification_loss
    173
    174            running_loss += loss.item()
    175            running_regression_loss += regression_loss.item()
    176            running_classification_loss += classification_loss.item()
    177        return running_loss / num, running_regression_loss / num, running_classification_loss / num