def __init__(self, cfg):
super(GeneralizedRCNN, self).__init__()
#self.backbone返回的是一个nn.Sequential的model,其中FPN出来的就是那多层特征
self.backbone = build_backbone(
cfg) #backbone 基础网络 提取特征 网络输出为out_channels = 256*4
#这里构造的是fcos_head
self.rpn = build_rpn(cfg, self.backbone.out_channels
) #生成提议的box以及对应的分类、位置回归以及中心度 对应fcos中的head
self.roi_heads = build_roi_heads(cfg, self.backbone.out_channels) #
def __init__(
self,
cfg,
confidence_thresholds_for_classes,
show_mask_heatmaps=False,
masks_per_dim=2,
min_image_size=224,
):
# self.CATEGORIES = ['__background', 'Person', 'Car', 'Train', 'Rider', 'Truck', 'Motorcycle', 'Bicycle', 'Bus']
self.CATEGORIES = ['__background', 'Car']
self.confidence_thresholds_for_classes = torch.tensor(
confidence_thresholds_for_classes)
self.cfg = cfg.clone()
self.min_image_size = min_image_size
self.model = {}
self.device = torch.device(cfg.MODEL.DEVICE)
backbone = build_backbone(cfg).to(self.device).eval()
fcos = build_rpn(cfg, backbone.out_channels).to(self.device).eval()
self.model["backbone"] = backbone
self.model["fcos"] = fcos
save_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg,
self.model,
save_dir=save_dir)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
self.transforms = self.build_transform()
mask_threshold = -1 if show_mask_heatmaps else 0.5
self.masker = Masker(threshold=mask_threshold, padding=1)
# used to make colors for each class
self.palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1])
self.cpu_device = torch.device("cpu")
self.confidence_thresholds_for_classes = torch.tensor(
confidence_thresholds_for_classes)
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
def train(cfg, local_rank, distributed):
writer = SummaryWriter('runs/{}'.format(cfg.OUTPUT_DIR))
##########################################################################
############################# Initial Model ##############################
##########################################################################
model = {}
device = torch.device(cfg.MODEL.DEVICE)
backbone = build_backbone(cfg).to(device)
fcos = build_rpn(cfg, backbone.out_channels).to(device)
if cfg.MODEL.ADV.USE_DIS_GLOBAL:
if cfg.MODEL.ADV.USE_DIS_P7:
dis_P7 = FCOSDiscriminator(
num_convs=cfg.MODEL.ADV.DIS_P7_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.GRL_WEIGHT_P7,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN).to(device)
if cfg.MODEL.ADV.USE_DIS_P6:
dis_P6 = FCOSDiscriminator(
num_convs=cfg.MODEL.ADV.DIS_P6_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.GRL_WEIGHT_P6,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN).to(device)
if cfg.MODEL.ADV.USE_DIS_P5:
dis_P5 = FCOSDiscriminator(
num_convs=cfg.MODEL.ADV.DIS_P5_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.GRL_WEIGHT_P5,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN).to(device)
if cfg.MODEL.ADV.USE_DIS_P4:
dis_P4 = FCOSDiscriminator(
num_convs=cfg.MODEL.ADV.DIS_P4_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.GRL_WEIGHT_P4,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN).to(device)
if cfg.MODEL.ADV.USE_DIS_P3:
dis_P3 = FCOSDiscriminator(
num_convs=cfg.MODEL.ADV.DIS_P3_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.GRL_WEIGHT_P3,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN).to(device)
if cfg.MODEL.ADV.USE_DIS_CENTER_AWARE:
if cfg.MODEL.ADV.USE_DIS_P7:
dis_P7_CA = FCOSDiscriminator_CA(
num_convs=cfg.MODEL.ADV.CA_DIS_P7_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.CA_GRL_WEIGHT_P7,
center_aware_weight=cfg.MODEL.ADV.CENTER_AWARE_WEIGHT,
center_aware_type=cfg.MODEL.ADV.CENTER_AWARE_TYPE,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN).to(device)
if cfg.MODEL.ADV.USE_DIS_P6:
dis_P6_CA = FCOSDiscriminator_CA(
num_convs=cfg.MODEL.ADV.CA_DIS_P6_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.CA_GRL_WEIGHT_P6,
center_aware_weight=cfg.MODEL.ADV.CENTER_AWARE_WEIGHT,
center_aware_type=cfg.MODEL.ADV.CENTER_AWARE_TYPE,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN).to(device)
if cfg.MODEL.ADV.USE_DIS_P5:
dis_P5_CA = FCOSDiscriminator_CA(
num_convs=cfg.MODEL.ADV.CA_DIS_P5_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.CA_GRL_WEIGHT_P5,
center_aware_weight=cfg.MODEL.ADV.CENTER_AWARE_WEIGHT,
center_aware_type=cfg.MODEL.ADV.CENTER_AWARE_TYPE,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN).to(device)
if cfg.MODEL.ADV.USE_DIS_P4:
dis_P4_CA = FCOSDiscriminator_CA(
num_convs=cfg.MODEL.ADV.CA_DIS_P4_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.CA_GRL_WEIGHT_P4,
center_aware_weight=cfg.MODEL.ADV.CENTER_AWARE_WEIGHT,
center_aware_type=cfg.MODEL.ADV.CENTER_AWARE_TYPE,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN).to(device)
if cfg.MODEL.ADV.USE_DIS_P3:
dis_P3_CA = FCOSDiscriminator_CA(
num_convs=cfg.MODEL.ADV.CA_DIS_P3_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.CA_GRL_WEIGHT_P3,
center_aware_weight=cfg.MODEL.ADV.CENTER_AWARE_WEIGHT,
center_aware_type=cfg.MODEL.ADV.CENTER_AWARE_TYPE,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN).to(device)
if cfg.MODEL.ADV.USE_DIS_CONDITIONAL:
if cfg.MODEL.ADV.USE_DIS_P7:
dis_P7_Cond = FCOSDiscriminator_CondA(
num_convs=cfg.MODEL.ADV.COND_DIS_P7_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.COND_GRL_WEIGHT_P7,
center_aware_weight=cfg.MODEL.ADV.CENTER_AWARE_WEIGHT,
# center_aware_type=cfg.MODEL.ADV.CENTER_AWARE_TYPE,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN,
class_align=cfg.MODEL.ADV.COND_CLASS,
reg_left_align=cfg.MODEL.ADV.COND_REG.LEFT,
reg_top_align=cfg.MODEL.ADV.COND_REG.TOP,
expand_dim=cfg.MODEL.ADV.COND_EXPAND).to(device)
if cfg.MODEL.ADV.USE_DIS_P6:
dis_P6_Cond = FCOSDiscriminator_CondA(
num_convs=cfg.MODEL.ADV.COND_DIS_P6_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.COND_GRL_WEIGHT_P6,
center_aware_weight=cfg.MODEL.ADV.CENTER_AWARE_WEIGHT,
# center_aware_type=cfg.MODEL.ADV.CENTER_AWARE_TYPE,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN,
class_align=cfg.MODEL.ADV.COND_CLASS,
reg_left_align=cfg.MODEL.ADV.COND_REG.LEFT,
reg_top_align=cfg.MODEL.ADV.COND_REG.TOP,
expand_dim=cfg.MODEL.ADV.COND_EXPAND).to(device)
if cfg.MODEL.ADV.USE_DIS_P5:
dis_P5_Cond = FCOSDiscriminator_CondA(
num_convs=cfg.MODEL.ADV.COND_DIS_P5_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.COND_GRL_WEIGHT_P5,
center_aware_weight=cfg.MODEL.ADV.CENTER_AWARE_WEIGHT,
# center_aware_type=cfg.MODEL.ADV.CENTER_AWARE_TYPE,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN,
class_align=cfg.MODEL.ADV.COND_CLASS,
reg_left_align=cfg.MODEL.ADV.COND_REG.LEFT,
reg_top_align=cfg.MODEL.ADV.COND_REG.TOP,
expand_dim=cfg.MODEL.ADV.COND_EXPAND).to(device)
if cfg.MODEL.ADV.USE_DIS_P4:
dis_P4_Cond = FCOSDiscriminator_CondA(
num_convs=cfg.MODEL.ADV.COND_DIS_P4_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.COND_GRL_WEIGHT_P4,
center_aware_weight=cfg.MODEL.ADV.CENTER_AWARE_WEIGHT,
# center_aware_type=cfg.MODEL.ADV.CENTER_AWARE_TYPE,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN,
class_align=cfg.MODEL.ADV.COND_CLASS,
reg_left_align=cfg.MODEL.ADV.COND_REG.LEFT,
reg_top_align=cfg.MODEL.ADV.COND_REG.TOP,
expand_dim=cfg.MODEL.ADV.COND_EXPAND).to(device)
if cfg.MODEL.ADV.USE_DIS_P3:
dis_P3_Cond = FCOSDiscriminator_CondA(
num_convs=cfg.MODEL.ADV.COND_DIS_P3_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.COND_GRL_WEIGHT_P3,
center_aware_weight=cfg.MODEL.ADV.CENTER_AWARE_WEIGHT,
# center_aware_type=cfg.MODEL.ADV.CENTER_AWARE_TYPE,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN,
class_align=cfg.MODEL.ADV.COND_CLASS,
reg_left_align=cfg.MODEL.ADV.COND_REG.LEFT,
reg_top_align=cfg.MODEL.ADV.COND_REG.TOP,
expand_dim=cfg.MODEL.ADV.COND_EXPAND).to(device)
if cfg.MODEL.ADV.USE_DIS_HEAD:
if cfg.MODEL.ADV.USE_DIS_P7:
dis_P7_HA = FCOSDiscriminator_HA(
num_convs=cfg.MODEL.ADV.HA_DIS_P7_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.HA_GRL_WEIGHT_P7,
center_aware_weight=cfg.MODEL.ADV.CENTER_AWARE_WEIGHT,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN).to(device)
if cfg.MODEL.ADV.USE_DIS_P6:
dis_P6_HA = FCOSDiscriminator_HA(
num_convs=cfg.MODEL.ADV.HA_DIS_P6_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.HA_GRL_WEIGHT_P6,
center_aware_weight=cfg.MODEL.ADV.CENTER_AWARE_WEIGHT,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN).to(device)
if cfg.MODEL.ADV.USE_DIS_P5:
dis_P5_HA = FCOSDiscriminator_HA(
num_convs=cfg.MODEL.ADV.HA_DIS_P5_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.HA_GRL_WEIGHT_P5,
center_aware_weight=cfg.MODEL.ADV.CENTER_AWARE_WEIGHT,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN).to(device)
if cfg.MODEL.ADV.USE_DIS_P4:
dis_P4_HA = FCOSDiscriminator_HA(
num_convs=cfg.MODEL.ADV.HA_DIS_P4_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.HA_GRL_WEIGHT_P4,
center_aware_weight=cfg.MODEL.ADV.CENTER_AWARE_WEIGHT,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN).to(device)
if cfg.MODEL.ADV.USE_DIS_P3:
dis_P3_HA = FCOSDiscriminator_HA(
num_convs=cfg.MODEL.ADV.HA_DIS_P3_NUM_CONVS,
grad_reverse_lambda=cfg.MODEL.ADV.HA_GRL_WEIGHT_P3,
center_aware_weight=cfg.MODEL.ADV.CENTER_AWARE_WEIGHT,
grl_applied_domain=cfg.MODEL.ADV.GRL_APPLIED_DOMAIN).to(device)
if cfg.MODEL.USE_SYNCBN:
assert is_pytorch_1_1_0_or_later(), \
"SyncBatchNorm is only available in pytorch >= 1.1.0"
backbone = torch.nn.SyncBatchNorm.convert_sync_batchnorm(backbone)
fcos = torch.nn.SyncBatchNorm.convert_sync_batchnorm(fcos)
if cfg.MODEL.ADV.USE_DIS_GLOBAL:
if cfg.MODEL.ADV.USE_DIS_P7:
dis_P7 = torch.nn.SyncBatchNorm.convert_sync_batchnorm(dis_P7)
if cfg.MODEL.ADV.USE_DIS_P6:
dis_P6 = torch.nn.SyncBatchNorm.convert_sync_batchnorm(dis_P6)
if cfg.MODEL.ADV.USE_DIS_P5:
dis_P5 = torch.nn.SyncBatchNorm.convert_sync_batchnorm(dis_P5)
if cfg.MODEL.ADV.USE_DIS_P4:
dis_P4 = torch.nn.SyncBatchNorm.convert_sync_batchnorm(dis_P4)
if cfg.MODEL.ADV.USE_DIS_P3:
dis_P3 = torch.nn.SyncBatchNorm.convert_sync_batchnorm(dis_P3)
if cfg.MODEL.ADV.USE_DIS_CENTER_AWARE:
if cfg.MODEL.ADV.USE_DIS_P7:
dis_P7_CA = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
dis_P7_CA)
if cfg.MODEL.ADV.USE_DIS_P6:
dis_P6_CA = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
dis_P6_CA)
if cfg.MODEL.ADV.USE_DIS_P5:
dis_P5_CA = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
dis_P5_CA)
if cfg.MODEL.ADV.USE_DIS_P4:
dis_P4_CA = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
dis_P4_CA)
if cfg.MODEL.ADV.USE_DIS_P3:
dis_P3_CA = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
dis_P3_CA)
if cfg.MODEL.ADV.USE_DIS_CONDITIONAL:
if cfg.MODEL.ADV.USE_DIS_P7:
dis_P7_Cond = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
dis_P7_Cond)
if cfg.MODEL.ADV.USE_DIS_P6:
dis_P6_Cond = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
dis_P6_Cond)
if cfg.MODEL.ADV.USE_DIS_P5:
dis_P5_Cond = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
dis_P5_Cond)
if cfg.MODEL.ADV.USE_DIS_P4:
dis_P4_Cond = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
dis_P4_Cond)
if cfg.MODEL.ADV.USE_DIS_P3:
dis_P3_Cond = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
dis_P3_Cond)
if cfg.MODEL.ADV.USE_DIS_HEAD:
if cfg.MODEL.ADV.USE_DIS_P7:
dis_P7_HA = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
dis_P7_HA)
if cfg.MODEL.ADV.USE_DIS_P6:
dis_P6_HA = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
dis_P6_HA)
if cfg.MODEL.ADV.USE_DIS_P5:
dis_P5_HA = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
dis_P5_HA)
if cfg.MODEL.ADV.USE_DIS_P4:
dis_P4_HA = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
dis_P4_HA)
if cfg.MODEL.ADV.USE_DIS_P3:
dis_P3_HA = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
dis_P3_HA)
##########################################################################
#################### Initial Optimizer and Scheduler #####################
##########################################################################
optimizer = {}
optimizer["backbone"] = make_optimizer(cfg, backbone, name='backbone')
optimizer["fcos"] = make_optimizer(cfg, fcos, name='fcos')
if cfg.MODEL.ADV.USE_DIS_GLOBAL:
if cfg.MODEL.ADV.USE_DIS_P7:
optimizer["dis_P7"] = make_optimizer(cfg,
dis_P7,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P6:
optimizer["dis_P6"] = make_optimizer(cfg,
dis_P6,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P5:
optimizer["dis_P5"] = make_optimizer(cfg,
dis_P5,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P4:
optimizer["dis_P4"] = make_optimizer(cfg,
dis_P4,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P3:
optimizer["dis_P3"] = make_optimizer(cfg,
dis_P3,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_CENTER_AWARE:
if cfg.MODEL.ADV.USE_DIS_P7:
optimizer["dis_P7_CA"] = make_optimizer(cfg,
dis_P7_CA,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P6:
optimizer["dis_P6_CA"] = make_optimizer(cfg,
dis_P6_CA,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P5:
optimizer["dis_P5_CA"] = make_optimizer(cfg,
dis_P5_CA,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P4:
optimizer["dis_P4_CA"] = make_optimizer(cfg,
dis_P4_CA,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P3:
optimizer["dis_P3_CA"] = make_optimizer(cfg,
dis_P3_CA,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_CONDITIONAL:
if cfg.MODEL.ADV.USE_DIS_P7:
optimizer["dis_P7_Cond"] = make_optimizer(cfg,
dis_P7_Cond,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P6:
optimizer["dis_P6_Cond"] = make_optimizer(cfg,
dis_P6_Cond,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P5:
optimizer["dis_P5_Cond"] = make_optimizer(cfg,
dis_P5_Cond,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P4:
optimizer["dis_P4_Cond"] = make_optimizer(cfg,
dis_P4_Cond,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P3:
optimizer["dis_P3_Cond"] = make_optimizer(cfg,
dis_P3_Cond,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_HEAD:
if cfg.MODEL.ADV.USE_DIS_P7:
optimizer["dis_P7_HA"] = make_optimizer(cfg,
dis_P7_HA,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P6:
optimizer["dis_P6_HA"] = make_optimizer(cfg,
dis_P6_HA,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P5:
optimizer["dis_P5_HA"] = make_optimizer(cfg,
dis_P5_HA,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P4:
optimizer["dis_P4_HA"] = make_optimizer(cfg,
dis_P4_HA,
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P3:
optimizer["dis_P3_HA"] = make_optimizer(cfg,
dis_P3_HA,
name='discriminator')
scheduler = {}
scheduler["backbone"] = make_lr_scheduler(cfg,
optimizer["backbone"],
name='backbone')
scheduler["fcos"] = make_lr_scheduler(cfg, optimizer["fcos"], name='fcos')
if cfg.MODEL.ADV.USE_DIS_GLOBAL:
if cfg.MODEL.ADV.USE_DIS_P7:
scheduler["dis_P7"] = make_lr_scheduler(cfg,
optimizer["dis_P7"],
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P6:
scheduler["dis_P6"] = make_lr_scheduler(cfg,
optimizer["dis_P6"],
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P5:
scheduler["dis_P5"] = make_lr_scheduler(cfg,
optimizer["dis_P5"],
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P4:
scheduler["dis_P4"] = make_lr_scheduler(cfg,
optimizer["dis_P4"],
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P3:
scheduler["dis_P3"] = make_lr_scheduler(cfg,
optimizer["dis_P3"],
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_CENTER_AWARE:
if cfg.MODEL.ADV.USE_DIS_P7:
scheduler["dis_P7_CA"] = make_lr_scheduler(cfg,
optimizer["dis_P7_CA"],
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P6:
scheduler["dis_P6_CA"] = make_lr_scheduler(cfg,
optimizer["dis_P6_CA"],
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P5:
scheduler["dis_P5_CA"] = make_lr_scheduler(cfg,
optimizer["dis_P5_CA"],
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P4:
scheduler["dis_P4_CA"] = make_lr_scheduler(cfg,
optimizer["dis_P4_CA"],
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P3:
scheduler["dis_P3_CA"] = make_lr_scheduler(cfg,
optimizer["dis_P3_CA"],
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_CONDITIONAL:
if cfg.MODEL.ADV.USE_DIS_P7:
scheduler["dis_P7_Cond"] = make_lr_scheduler(
cfg, optimizer["dis_P7_Cond"], name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P6:
scheduler["dis_P6_Cond"] = make_lr_scheduler(
cfg, optimizer["dis_P6_Cond"], name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P5:
scheduler["dis_P5_Cond"] = make_lr_scheduler(
cfg, optimizer["dis_P5_Cond"], name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P4:
scheduler["dis_P4_Cond"] = make_lr_scheduler(
cfg, optimizer["dis_P4_Cond"], name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P3:
scheduler["dis_P3_Cond"] = make_lr_scheduler(
cfg, optimizer["dis_P3_Cond"], name='discriminator')
if cfg.MODEL.ADV.USE_DIS_HEAD:
if cfg.MODEL.ADV.USE_DIS_P7:
scheduler["dis_P7_HA"] = make_lr_scheduler(cfg,
optimizer["dis_P7_HA"],
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P6:
scheduler["dis_P6_HA"] = make_lr_scheduler(cfg,
optimizer["dis_P6_HA"],
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P5:
scheduler["dis_P5_HA"] = make_lr_scheduler(cfg,
optimizer["dis_P5_HA"],
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P4:
scheduler["dis_P4_HA"] = make_lr_scheduler(cfg,
optimizer["dis_P4_HA"],
name='discriminator')
if cfg.MODEL.ADV.USE_DIS_P3:
scheduler["dis_P3_HA"] = make_lr_scheduler(cfg,
optimizer["dis_P3_HA"],
name='discriminator')
##########################################################################
######################## DistributedDataParallel #########################
##########################################################################
if distributed:
backbone = torch.nn.parallel.DistributedDataParallel(
backbone,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
fcos = torch.nn.parallel.DistributedDataParallel(
fcos,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_GLOBAL:
if cfg.MODEL.ADV.USE_DIS_P7:
dis_P7 = torch.nn.parallel.DistributedDataParallel(
dis_P7,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_P6:
dis_P6 = torch.nn.parallel.DistributedDataParallel(
dis_P6,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_P5:
dis_P5 = torch.nn.parallel.DistributedDataParallel(
dis_P5,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_P4:
dis_P4 = torch.nn.parallel.DistributedDataParallel(
dis_P4,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_P3:
dis_P3 = torch.nn.parallel.DistributedDataParallel(
dis_P3,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_CENTER_AWARE:
if cfg.MODEL.ADV.USE_DIS_P7:
dis_P7_CA = torch.nn.parallel.DistributedDataParallel(
dis_P7_CA,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_P6:
dis_P6_CA = torch.nn.parallel.DistributedDataParallel(
dis_P6_CA,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_P5:
dis_P5_CA = torch.nn.parallel.DistributedDataParallel(
dis_P5_CA,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_P4:
dis_P4_CA = torch.nn.parallel.DistributedDataParallel(
dis_P4_CA,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_P3:
dis_P3_CA = torch.nn.parallel.DistributedDataParallel(
dis_P3_CA,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_CONDITIONAL:
if cfg.MODEL.ADV.USE_DIS_P7:
dis_P7_Cond = torch.nn.parallel.DistributedDataParallel(
dis_P7_Cond,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_P6:
dis_P6_Cond = torch.nn.parallel.DistributedDataParallel(
dis_P6_Cond,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_P5:
dis_P5_Cond = torch.nn.parallel.DistributedDataParallel(
dis_P5_Cond,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_P4:
dis_P4_Cond = torch.nn.parallel.DistributedDataParallel(
dis_P4_Cond,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_P3:
dis_P3_Cond = torch.nn.parallel.DistributedDataParallel(
dis_P3_Cond,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_HEAD:
if cfg.MODEL.ADV.USE_DIS_P7:
dis_P7_HA = torch.nn.parallel.DistributedDataParallel(
dis_P7_HA,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_P6:
dis_P6_HA = torch.nn.parallel.DistributedDataParallel(
dis_P6_HA,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_P5:
dis_P5_HA = torch.nn.parallel.DistributedDataParallel(
dis_P5_HA,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_P4:
dis_P4_HA = torch.nn.parallel.DistributedDataParallel(
dis_P4_HA,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
if cfg.MODEL.ADV.USE_DIS_P3:
dis_P3_HA = torch.nn.parallel.DistributedDataParallel(
dis_P3_HA,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False)
##########################################################################
########################### Save Model to Dict ###########################
##########################################################################
model["backbone"] = backbone
model["fcos"] = fcos
if cfg.MODEL.ADV.USE_DIS_GLOBAL:
if cfg.MODEL.ADV.USE_DIS_P7:
model["dis_P7"] = dis_P7
if cfg.MODEL.ADV.USE_DIS_P6:
model["dis_P6"] = dis_P6
if cfg.MODEL.ADV.USE_DIS_P5:
model["dis_P5"] = dis_P5
if cfg.MODEL.ADV.USE_DIS_P4:
model["dis_P4"] = dis_P4
if cfg.MODEL.ADV.USE_DIS_P3:
model["dis_P3"] = dis_P3
if cfg.MODEL.ADV.USE_DIS_CENTER_AWARE:
if cfg.MODEL.ADV.USE_DIS_P7:
model["dis_P7_CA"] = dis_P7_CA
if cfg.MODEL.ADV.USE_DIS_P6:
model["dis_P6_CA"] = dis_P6_CA
if cfg.MODEL.ADV.USE_DIS_P5:
model["dis_P5_CA"] = dis_P5_CA
if cfg.MODEL.ADV.USE_DIS_P4:
model["dis_P4_CA"] = dis_P4_CA
if cfg.MODEL.ADV.USE_DIS_P3:
model["dis_P3_CA"] = dis_P3_CA
if cfg.MODEL.ADV.USE_DIS_CONDITIONAL:
if cfg.MODEL.ADV.USE_DIS_P7:
model["dis_P7_Cond"] = dis_P7_Cond
if cfg.MODEL.ADV.USE_DIS_P6:
model["dis_P6_Cond"] = dis_P6_Cond
if cfg.MODEL.ADV.USE_DIS_P5:
model["dis_P5_Cond"] = dis_P5_Cond
if cfg.MODEL.ADV.USE_DIS_P4:
model["dis_P4_Cond"] = dis_P4_Cond
if cfg.MODEL.ADV.USE_DIS_P3:
model["dis_P3_Cond"] = dis_P3_Cond
if cfg.MODEL.ADV.USE_DIS_HEAD:
if cfg.MODEL.ADV.USE_DIS_P7:
model["dis_P7_HA"] = dis_P7_HA
if cfg.MODEL.ADV.USE_DIS_P6:
model["dis_P6_HA"] = dis_P6_HA
if cfg.MODEL.ADV.USE_DIS_P5:
model["dis_P5_HA"] = dis_P5_HA
if cfg.MODEL.ADV.USE_DIS_P4:
model["dis_P4_HA"] = dis_P4_HA
if cfg.MODEL.ADV.USE_DIS_P3:
model["dis_P3_HA"] = dis_P3_HA
##########################################################################
################################ Training ################################
##########################################################################
arguments = {}
arguments["iteration"] = 0
arguments["use_dis_global"] = cfg.MODEL.ADV.USE_DIS_GLOBAL
arguments["use_dis_ca"] = cfg.MODEL.ADV.USE_DIS_CENTER_AWARE
arguments["use_dis_conditional"] = cfg.MODEL.ADV.USE_DIS_CONDITIONAL
arguments["use_dis_ha"] = cfg.MODEL.ADV.USE_DIS_HEAD
arguments["ga_dis_lambda"] = cfg.MODEL.ADV.GA_DIS_LAMBDA
arguments["ca_dis_lambda"] = cfg.MODEL.ADV.CA_DIS_LAMBDA
arguments["cond_dis_lambda"] = cfg.MODEL.ADV.COND_DIS_LAMBDA
arguments["ha_dis_lambda"] = cfg.MODEL.ADV.HA_DIS_LAMBDA
arguments["use_feature_layers"] = []
if cfg.MODEL.ADV.USE_DIS_P7:
arguments["use_feature_layers"].append("P7")
if cfg.MODEL.ADV.USE_DIS_P6:
arguments["use_feature_layers"].append("P6")
if cfg.MODEL.ADV.USE_DIS_P5:
arguments["use_feature_layers"].append("P5")
if cfg.MODEL.ADV.USE_DIS_P4:
arguments["use_feature_layers"].append("P4")
if cfg.MODEL.ADV.USE_DIS_P3:
arguments["use_feature_layers"].append("P3")
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(cfg, model, optimizer, scheduler,
output_dir, save_to_disk)
extra_checkpoint_data = checkpointer.load(f=cfg.MODEL.WEIGHT,
load_dis=True,
load_opt_sch=False)
# arguments.update(extra_checkpoint_data)
# Initial dataloader (both target and source domain)
data_loader = {}
data_loader["source"] = make_data_loader_source(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
data_loader["target"] = make_data_loader_target(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
do_train(model, data_loader, optimizer, scheduler, checkpointer, device,
checkpoint_period, arguments, cfg, run_test, distributed, writer)
return model
def main():
parser = argparse.ArgumentParser(description="PyTorch Object Detection Inference")
parser.add_argument(
"--config-file",
default="/private/home/fmassa/github/detectron.pytorch_v2/configs/e2e_faster_rcnn_R_50_C4_1x_caffe2.yaml",
metavar="FILE",
help="path to config file",
)
parser.add_argument("--local_rank", type=int, default=0)
parser.add_argument(
"opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
distributed = num_gpus > 1
if distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(
backend="nccl", init_method="env://"
)
synchronize()
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
save_dir = ""
logger = setup_logger("fcos_core", save_dir, get_rank())
logger.info("Using {} GPUs".format(num_gpus))
logger.info(cfg)
logger.info("Collecting env info (might take some time)")
logger.info("\n" + collect_env_info())
# model = build_detection_model(cfg)
# model.to(cfg.MODEL.DEVICE)
from fcos_core.modeling.backbone import build_backbone
from fcos_core.modeling.rpn.rpn import build_rpn
model = {}
model["backbone"] = build_backbone(cfg).to(cfg.MODEL.DEVICE)
model["fcos"] = build_rpn(cfg, model["backbone"].out_channels).to(cfg.MODEL.DEVICE)
output_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg, model, save_dir=output_dir)
_ = checkpointer.load(f=cfg.MODEL.WEIGHT, load_dis=False, load_opt_sch=False)
iou_types = ("bbox",)
if cfg.MODEL.MASK_ON:
iou_types = iou_types + ("segm",)
if cfg.MODEL.KEYPOINT_ON:
iou_types = iou_types + ("keypoints",)
output_folders = [None] * len(cfg.DATASETS.TEST)
dataset_names = cfg.DATASETS.TEST
if cfg.OUTPUT_DIR:
for idx, dataset_name in enumerate(dataset_names):
output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name)
mkdir(output_folder)
output_folders[idx] = output_folder
data_loaders_val = make_data_loader(cfg, is_train=False, is_distributed=distributed)
for output_folder, dataset_name, data_loader_val in zip(output_folders, dataset_names, data_loaders_val):
inference(
model,
data_loader_val,
dataset_name=dataset_name,
iou_types=iou_types,
box_only=False if cfg.MODEL.FCOS_ON or cfg.MODEL.RETINANET_ON else cfg.MODEL.RPN_ONLY,
device=cfg.MODEL.DEVICE,
expected_results=cfg.TEST.EXPECTED_RESULTS,
expected_results_sigma_tol=cfg.TEST.EXPECTED_RESULTS_SIGMA_TOL,
output_folder=output_folder,
)
synchronize()
def train(cfg, local_rank, distributed):
writer = SummaryWriter('runs/{}'.format(cfg.OUTPUT_DIR))
##########################################################################
############################# Initial Model ##############################
##########################################################################
model = {}
device = torch.device(cfg.MODEL.DEVICE)
backbone = build_backbone(cfg).to(device)
fcos = build_rpn(cfg, backbone.out_channels).to(device)
if cfg.MODEL.USE_SYNCBN:
assert is_pytorch_1_1_0_or_later(), \
"SyncBatchNorm is only available in pytorch >= 1.1.0"
backbone = torch.nn.SyncBatchNorm.convert_sync_batchnorm(backbone)
fcos = torch.nn.SyncBatchNorm.convert_sync_batchnorm(fcos)
##########################################################################
#################### Initial Optimizer and Scheduler #####################
##########################################################################
optimizer = {}
optimizer["backbone"] = make_optimizer(cfg, backbone, name='backbone')
optimizer["fcos"] = make_optimizer(cfg, fcos, name='fcos')
scheduler = {}
scheduler["backbone"] = make_lr_scheduler(cfg, optimizer["backbone"], name='backbone')
scheduler["fcos"] = make_lr_scheduler(cfg, optimizer["fcos"], name='fcos')
if distributed:
backbone = torch.nn.parallel.DistributedDataParallel(
backbone, device_ids=[local_rank], output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False
)
fcos = torch.nn.parallel.DistributedDataParallel(
fcos, device_ids=[local_rank], output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False
)
########################### Save Model to Dict ###########################
##########################################################################
model["backbone"] = backbone
model["fcos"] = fcos
##########################################################################
################################ Training ################################
##########################################################################
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(
cfg, model, optimizer, scheduler, output_dir, save_to_disk
)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
data_loader = make_data_loader_source(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
do_train_base(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
cfg,
run_test,
distributed,
writer
)
return model