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