def __init__(self, cfg, vgg: nn.ModuleList, extras: nn.ModuleList,
classification_headers: nn.ModuleList,
regression_headers: nn.ModuleList):
"""Compose a SSD model using the given components.
"""
super(SSD, self).__init__()
self.cfg = cfg
self.num_classes = cfg.MODEL.NUM_CLASSES
self.vgg = vgg
self.extras = extras
self.classification_headers = classification_headers
self.regression_headers = regression_headers
self.l2_norm = L2Norm(512, scale=20)
self.criterion = MultiBoxLoss(neg_pos_ratio=cfg.MODEL.NEG_POS_RATIO)
self.priors = None
if cfg.MODEL.SELF_SUPERVISED:
n_out = 256
if cfg.MODEL.SELF_SUPERVISOR.TYPE == "jigen":
self.ss_classifier = nn.Linear(
n_out * 3 * 3, cfg.MODEL.SELF_SUPERVISOR.CLASSES + 1)
else:
self.ss_dropout = nn.Dropout(
p=cfg.MODEL.SELF_SUPERVISOR.DROPOUT)
self.ss_classifier = nn.Linear(n_out * 1 * 1, 4)
self.reset_parameters()
def train(cfg, args):
logger = logging.getLogger('SSD.trainer')
# -----------------------------------------------------------------------------
# Model
# -----------------------------------------------------------------------------
model = build_mobilev1_ssd_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
if args.resume:
logger.info("Resume from the model {}".format(args.resume))
model.load(args.resume)
else:
logger.info("Init from base net {}".format(args.vgg))
model.init_from_base_net(args.vgg)
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)
# -----------------------------------------------------------------------------
# Optimizer
# -----------------------------------------------------------------------------
lr = cfg.SOLVER.LR * args.num_gpus # scale by num gpus
optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=cfg.SOLVER.MOMENTUM, weight_decay=cfg.SOLVER.WEIGHT_DECAY)
# -----------------------------------------------------------------------------
# Criterion
# -----------------------------------------------------------------------------
criterion = MultiBoxLoss(neg_pos_ratio=cfg.MODEL.NEG_POS_RATIO)
# -----------------------------------------------------------------------------
# Scheduler
# -----------------------------------------------------------------------------
milestones = [step // args.num_gpus for step in cfg.SOLVER.LR_STEPS]
scheduler = WarmupMultiStepLR(optimizer=optimizer,
milestones=milestones,
gamma=cfg.SOLVER.GAMMA,
warmup_factor=cfg.SOLVER.WARMUP_FACTOR,
warmup_iters=cfg.SOLVER.WARMUP_ITERS)
# -----------------------------------------------------------------------------
# Dataset
# -----------------------------------------------------------------------------
train_transform = TrainAugmentation(cfg.INPUT.IMAGE_SIZE, cfg.INPUT.PIXEL_MEAN)
target_transform = MatchPrior(PriorBox(cfg)(), cfg.MODEL.CENTER_VARIANCE, cfg.MODEL.SIZE_VARIANCE, cfg.MODEL.THRESHOLD)
train_dataset = build_dataset(dataset_list=cfg.DATASETS.TRAIN, transform=train_transform, target_transform=target_transform)
logger.info("Train dataset size: {}".format(len(train_dataset)))
if args.distributed:
sampler = torch.utils.data.DistributedSampler(train_dataset)
else:
sampler = torch.utils.data.RandomSampler(train_dataset)
batch_sampler = torch.utils.data.sampler.BatchSampler(sampler=sampler, batch_size=cfg.SOLVER.BATCH_SIZE, drop_last=False)
batch_sampler = samplers.IterationBasedBatchSampler(batch_sampler, num_iterations=cfg.SOLVER.MAX_ITER // args.num_gpus)
train_loader = DataLoader(train_dataset, num_workers=4, batch_sampler=batch_sampler)
return do_train(cfg, model, train_loader, optimizer, scheduler, criterion, device, args)
def __init__(self, cfg,
vgg: nn.ModuleList,
extras: nn.ModuleList,
classification_headers: nn.ModuleList,
regression_headers: nn.ModuleList):
"""Compose a SSD model using the given components.
"""
super(SSD, self).__init__()
self.cfg = cfg
self.num_classes = cfg.MODEL.NUM_CLASSES
self.vgg = vgg
self.extras = extras
self.classification_headers = classification_headers
self.regression_headers = regression_headers
self.l2_norm = L2Norm(512, scale=20)
self.criterion = MultiBoxLoss(neg_pos_ratio=cfg.MODEL.NEG_POS_RATIO)
self.priors = None
self.reset_parameters()
def __init__(self, cfg,
vgg: nn.ModuleList,
extras: nn.ModuleList,
classification_headers: nn.ModuleList,
regression_headers: nn.ModuleList,
downsample_layers_index:list):
"""Compose a SSD model using the given components.
"""
super(SSD, self).__init__()
self.cfg = cfg
self.num_classes = cfg.MODEL.NUM_CLASSES
self.vgg = vgg
self.extras = extras
self.classification_headers = classification_headers
self.regression_headers = regression_headers
self.l2_norm = L2Norm(512, scale=20)
self.criterion = MultiBoxLoss(neg_pos_ratio=cfg.MODEL.NEG_POS_RATIO)
self.priors = None
self.downsample_layers_index = downsample_layers_index
# FCN part
self.fcn_module = []
self.conv1 = nn.Conv2d(512, 512, 1)
self.fcn_module.append(self.conv1)
self.bn1 = nn.BatchNorm2d(512)
self.fcn_module.append(self.bn1)
self.relu1 = nn.ReLU()
self.conv2 = nn.Conv2d(256, 64, 3, padding=1)
self.fcn_module.append(self.conv2)
self.bn2 = nn.BatchNorm2d(64)
self.fcn_module.append(self.bn2)
self.relu2 = nn.ReLU()
self.conv3 = nn.Conv2d(64, 1, 1)
self.sigmoid = nn.Sigmoid()
self.unpool1 = nn.Upsample(scale_factor=2, mode='bilinear')
self.unpool1_conv2d = nn.Conv2d(1024,512,1)
self.fcn_module.append(self.unpool1_conv2d)
self.unpool2 = nn.Upsample(scale_factor=2, mode='bilinear')
self.unpool2_conv2d = nn.Conv2d(512,256,1)
self.fcn_module.append(self.unpool2_conv2d)
self.fcn_module = nn.ModuleList(self.fcn_module)
self.reset_parameters()