what.models.detection.ssd.mobilenet_v1_ssd API documentation

# what.models.detection.ssd.mobilenet_v1_ssd

View Source

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

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

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

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

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

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

    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