From 0e653f5f2c5067eabb1c8edc6095cdac24d3ba01 Mon Sep 17 00:00:00 2001 From: Hao Date: Fri, 19 Oct 2018 12:35:53 +1300 Subject: [PATCH] add squeezenet ssd lite --- models/EMPTY | 0 train_ssd.py | 12 +++-- vision/ssd/squeezenet_ssd_lite.py | 85 +++++++++++++++++++++++++++++++ 3 files changed, 93 insertions(+), 4 deletions(-) create mode 100644 models/EMPTY create mode 100644 vision/ssd/squeezenet_ssd_lite.py diff --git a/models/EMPTY b/models/EMPTY new file mode 100644 index 00000000..e69de29b diff --git a/train_ssd.py b/train_ssd.py index 4adb06f1..0ba970ea 100644 --- a/train_ssd.py +++ b/train_ssd.py @@ -13,7 +13,7 @@ from vision.ssd.vgg_ssd import create_vgg_ssd from vision.ssd.mobilenetv1_ssd import create_mobilenetv1_ssd from vision.ssd.mobilenetv1_ssd_lite import create_mobilenetv1_ssd_lite -from vision.ssd.fpn_mobilenetv1_ssd import create_fpn_mobilenetv1_ssd +from vision.ssd.squeezenet_ssd_lite import create_squeezenet_ssd_lite from vision.datasets.voc_dataset import VOCDataset from vision.datasets.open_images import OpenImagesDataset from vision.nn.multibox_loss import MultiboxLoss @@ -92,10 +92,14 @@ help='Directory for saving checkpoint models') +logging.basicConfig(stream=sys.stdout, level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') args = parser.parse_args() DEVICE = torch.device("cuda:0" if torch.cuda.is_available() and args.use_cuda else "cpu") + if args.use_cuda and torch.cuda.is_available(): torch.backends.cudnn.benchmark = True + logging.info("Use Cuda.") def train(loader, net, criterion, optimizer, device, debug_steps=100, epoch=-1): @@ -159,9 +163,6 @@ def test(loader, net, criterion, device): if __name__ == '__main__': - logging.basicConfig(stream=sys.stdout, level=logging.INFO, - format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') - timer = Timer() logging.info(args) @@ -174,6 +175,9 @@ def test(loader, net, criterion, device): elif args.net == 'mb1-ssd-lite': create_net = create_mobilenetv1_ssd_lite config = mobilenetv1_ssd_config + elif args.net == 'sq-ssd-lite': + create_net = create_squeezenet_ssd_lite + config = mobilenetv1_ssd_config else: logging.fatal("The net type is wrong.") parser.print_help(sys.stderr) diff --git a/vision/ssd/squeezenet_ssd_lite.py b/vision/ssd/squeezenet_ssd_lite.py new file mode 100644 index 00000000..fd0882cd --- /dev/null +++ b/vision/ssd/squeezenet_ssd_lite.py @@ -0,0 +1,85 @@ +import torch +from torch.nn import Conv2d, Sequential, ModuleList, ReLU, BatchNorm2d +from torchvision import models + +from .ssd import SSD +from .predictor import Predictor +from .config import mobilenetv1_ssd_config as config + + +def SeperableConv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0): + """Replace Conv2d with a depthwise Conv2d and Pointwise Conv2d. + """ + return Sequential( + Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size, + groups=in_channels, stride=stride, padding=padding), + ReLU(), + Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1), + ) + + +def create_squeezenet_ssd_lite(num_classes, is_test=False): + base_net = models.squeezenet1_1(False).features # disable dropout layer + + source_layer_indexes = [ + 12 + ] + extras = ModuleList([ + Sequential( + Conv2d(in_channels=512, out_channels=256, kernel_size=1), + ReLU(), + SeperableConv2d(in_channels=256, out_channels=512, kernel_size=3, stride=2, padding=1), + ), + Sequential( + Conv2d(in_channels=512, out_channels=256, kernel_size=1), + ReLU(), + SeperableConv2d(in_channels=256, out_channels=512, kernel_size=3, stride=2, padding=1), + ), + Sequential( + Conv2d(in_channels=512, out_channels=128, kernel_size=1), + ReLU(), + SeperableConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1), + ), + Sequential( + Conv2d(in_channels=256, out_channels=128, kernel_size=1), + ReLU(), + SeperableConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1), + ), + Sequential( + Conv2d(in_channels=256, out_channels=128, kernel_size=1), + ReLU(), + SeperableConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1) + ) + ]) + + regression_headers = ModuleList([ + SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1), + SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1), + SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1), + SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1), + SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1), + Conv2d(in_channels=256, out_channels=6 * 4, kernel_size=1), + ]) + + classification_headers = ModuleList([ + SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1), + SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1), + SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1), + SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1), + SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1), + Conv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=1), + ]) + + return SSD(num_classes, base_net, source_layer_indexes, + extras, classification_headers, regression_headers, is_test=is_test, config=config) + + +def create_squeezenet_ssd_lite_predictor(net, candidate_size=200, nms_method=None, sigma=0.5, device=torch.device('cpu')): + predictor = Predictor(net, config.image_size, config.image_mean, + config.image_std, + nms_method=nms_method, + iou_threshold=config.iou_threshold, + candidate_size=candidate_size, + sigma=sigma, + device=device) + return predictor \ No newline at end of file