what.models.detection.ssd.mobilenet_v2_ssd_lite
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:
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)
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
step( self, loader, criterion, optimizer, device, debug_steps=100, epoch=-1):
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/'):
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):
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