def _init_pretrained_model(self, num_classes): box_roi_pool = MultiScaleRoIAlign( featmap_names=[0, 1, 2, 3], # + "pool" -> 5 feature maps output_size=7, sampling_ratio=2) model = fasterrcnn_resnet50_fpn(pretrained=True, min_size=config.IMAGE_SIZE, box_nms_thresh=.5, box_roi_pool=box_roi_pool) torch.manual_seed(0) # Init the same params in all processes model.roi_heads.box_predictor = FastRCNNPredictor( in_channels=model.roi_heads.box_head.fc7.out_features, num_classes=num_classes) model.rpn.anchor_generator = AnchorGenerator( sizes=[[16], [32], [64], [128], [256]], aspect_ratios=[.25, .5, 1., 2., 4.]) model.rpn.head = RPNHead(in_channels=256, num_anchors=5) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model.to(device) model = DDP(model, find_unused_parameters=True) if self.is_master(): print(model) return model
def __init__(self, threshold=0.7, category_map=None, weights=None): if weights is None: self.pytorch_model = fasterrcnn_resnet50_fpn( pretrained=True, pretrained_backbone=True) else: self.pytorch_model = fasterrcnn_resnet50_fpn( pretrained=False, pretrained_backbone=False) self.pytorch_model.load_state_dict(torch.load(weights)) if category_map is None: category_map = CATEGORY_MAP self.category_map = category_map self.threshold = threshold self.img_numpy = None self.img_tensor = None self.img_out = None self.filter_objs = dict()
def model(num_classes: int, backbone: Optional[nn.Module] = None, remove_internal_transforms: bool = True, **faster_rcnn_kwargs) -> nn.Module: """ FasterRCNN model given by torchvision Args: num_classes (int): Number of classes. backbone (nn.Module): Backbone model to use. Defaults to a resnet50_fpn model. Return: nn.Module """ if backbone is None: model = fasterrcnn_resnet50_fpn(pretrained=True, **faster_rcnn_kwargs) in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor( in_features, num_classes) backbone_param_groups = resnet_fpn.param_groups(model.backbone) else: model = FasterRCNN(backbone, num_classes=num_classes, **faster_rcnn_kwargs) backbone_param_groups = backbone.param_groups() patch_param_groups(model=model, backbone_param_groups=backbone_param_groups) if remove_internal_transforms: remove_internal_model_transforms(model) return model
def main(): dataset = IcartoonDataset(transforms=True) indices = torch.randperm(len(dataset)).tolist() dataset = torch.utils.data.Subset(dataset, indices[:-50]) dataset_test = torch.utils.data.Subset(dataset, indices[-50:]) test_loader = torch.utils.data.DataLoader(dataset_test, num_workers=4, collate_fn=collate_fn) num_classes = 1 model = fasterrcnn_resnet50_fpn(pretrained=False, num_classes=num_classes) # 构造一个优化器 params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005) # 和学习率调度程序 lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1) epoch = 5 for e in range(epoch): train(model, device='cpu', train_loader=test_loader, optimizer=optimizer, epoch=e) lr_scheduler.step() torch.save(model.state_dict(), './{}.model'.format(e)) pass
def create_fastercnn( num_classes: int = 91, backbone: nn.Module = None, **kwargs, ): """ Creates Faster RCNN implementation based on torchvision library. Args: num_classes (int) : number of classes. Do not have class_id "0" it is reserved as background. num_classes = number of classes to label + 1 for background. """ if backbone is None: # Creates the default fasterrcnn as given in pytorch. Trained on COCO dataset model = fasterrcnn_resnet50_fpn( pretrained=True, **kwargs, ) in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor( in_features, num_classes) else: model = FasterRCNN(backbone, num_classes=num_classes, **kwargs) return model
def faster_rcnn_res50(self_pretrained, num_classes): if not self_pretrained: print('load res50 backbone pretrained on COCO') pretrained = False if self_pretrained else True # pretrain on COCO model = fasterrcnn_resnet50_fpn(pretrained=pretrained) in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_channels=in_features, num_classes=num_classes) return model
def get_pretrained_faster_rcnn(num_classes): """return nn.Module""" model = fasterrcnn_resnet50_fpn(pretrained=True) # get number of input channels for the classifier in_channels = model.roi_heads.box_predictor.cls_score.in_features # replace box predictor and add class 'background' model.roi_heads.box_predictor = FastRCNNPredictor(in_channels=in_channels, num_classes=num_classes + 1) return model
def get_fasterrcnn_resnet50_fpn(num_classes=2, pretrained=True, pretrained_backbone=True, trainable_backbone_layers=5): model = fasterrcnn_resnet50_fpn( pretrained=pretrained, trainable_backbone_layers=trainable_backbone_layers, pretrained_backbone=pretrained_backbone, ) in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) return model
def __init__(self): self.model = fasterrcnn_resnet50_fpn(pretrained=True) classes = 2 # one for the background in_features = self.model.roi_heads.box_predictor.cls_score.in_features self.model.roi_heads.box_predictor = FastRCNNPredictor( in_features, classes) checkpoint = torch.load('leaf_od_weights.pt', map_location=torch.device('cpu')) self.model.load_state_dict(checkpoint['model_state_dict']) self.model.eval()
def get_model_fastrcnn(num_classes, ckpt): model = fasterrcnn_resnet50_fpn(pretrained=True) # get number of input features for the classifier in_features = model.roi_heads.box_predictor.cls_score.in_features # replace the pre-trained head with a new one model.roi_heads.box_predictor = FastRCNNPredictor( in_features, Configuration.num_classes) if ckpt is not None: model.load_state_dict(torch.load(ckpt)) return model
def __init__(self, num_classes, rpn_pre_nms_top_n_train=100, rpn_pre_nms_top_n_test=100, rpn_post_nms_top_n_train=100, rpn_post_nms_top_n_test=100, rpn_anchor_generator=None): super().__init__() # load a model pre-trained pre-trained on COCO self.model = fasterrcnn_resnet50_fpn(pretrained=True, rpn_pre_nms_top_n_train=100, rpn_pre_nms_top_n_test=100, rpn_post_nms_top_n_train=100, rpn_post_nms_top_n_test=100, rpn_anchor_generator=rpn_anchor_generator) # get number of input features for the classifier in_features = self.model.roi_heads.box_predictor.cls_score.in_features # replace the pre-trained head with a new one self.model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
def fasterRCNN_ResNet50_fpn(num_classes=72, anchor_sizes=((32, ), (64, ), (128, ), (256, ), (512, )), aspect_ratios=((0.5, 1.0, 2.0), ) * 5, image_mean=None, image_std=None, min_size=200, max_size=600): """ Generate a Faster R-CNN with ResNet50 backbone :param num_classes: number of output classes of the model (including the background) :param anchor_sizes: List[Tuple[int],...], anchor sizes :param aspect_ratios: List[Tuple[int],...], aspect ratios For both anchor_sizes and aspect_ratios, you should specify as many tuples as the number of feature maps that you want to extract the RoIs from, anchor_sizes and aspect_ratios should have the same number of elements, aka. len(anchor_sizes) == len(aspect_ratios), e.g. sizes=((32,), (64,), (128,), (256,), (512,)), aspect_ratios=((0.5, 1.0, 2.0),(0.5, 1.0, 2.0),(0.5, 1.0, 2.0),(0.5, 1.0, 2.0),(0.5, 1.0, 2.0)), which will generate 1x3 anchors per spatial location with 1 size and 3 different ratios for 5 feature maps. ResNet 50 with fpn backbone outputs 5 feature maps so the length of anchor_sizes and the aspect_ratios must be 5. If input size is 200x200, the output feature maps are {layer1:50x50, layer2:25x25, layer3:13x13, layer4:7x7, pool:4x4}. :param image_mean: List[int,int,int], mean values for RGB 3 channels :param image_std: List[int,int,int], std. values for RGB 3 channels Mean and std. values are used for normalization, calculated from the training dataset :return: """ if len(anchor_sizes) != 5 or len(aspect_ratios) != 5: LOG.warning( "The length of anchor_sizes and the aspect_ratios must be 5!") raise ValueError( "The length of anchor_sizes and the aspect_ratios must be 5!") args = { 'rpn_anchor_generator': AnchorGenerator(sizes=anchor_sizes, aspect_ratios=aspect_ratios), 'image_mean': image_mean, 'image_std': image_std, 'min_size': min_size, 'max_size': max_size } model = fasterrcnn_resnet50_fpn(pretrained=False, num_classes=num_classes, pretrained_backbone=False, **args) return model
def __init__( self, learning_rate: float = 0.0001, num_classes: int = 91, backbone: Optional[str] = None, fpn: bool = True, pretrained: bool = False, pretrained_backbone: bool = True, trainable_backbone_layers: int = 3, **kwargs: Any, ): """ Args: learning_rate: the learning rate num_classes: number of detection classes (including background) backbone: Pretained backbone CNN architecture. fpn: If True, creates a Feature Pyramind Network on top of Resnet based CNNs. pretrained: if true, returns a model pre-trained on COCO train2017 pretrained_backbone: if true, returns a model with backbone pre-trained on Imagenet trainable_backbone_layers: number of trainable resnet layers starting from final block """ if not _TORCHVISION_AVAILABLE: # pragma: no cover raise ModuleNotFoundError( 'You want to use `torchvision` which is not installed yet.') super().__init__() self.learning_rate = learning_rate self.num_classes = num_classes self.backbone = backbone if backbone is None: self.model = fasterrcnn_resnet50_fpn( pretrained=pretrained, pretrained_backbone=pretrained_backbone, trainable_backbone_layers=trainable_backbone_layers, ) in_features = self.model.roi_heads.box_predictor.cls_score.in_features self.model.roi_heads.box_predictor = FastRCNNPredictor( in_features, self.num_classes) else: backbone_model = create_fasterrcnn_backbone( self.backbone, fpn, pretrained_backbone, trainable_backbone_layers, **kwargs, ) self.model = torchvision_FasterRCNN(backbone_model, num_classes=num_classes, **kwargs)
def main(): device = torch.device( 'cuda') if torch.cuda.is_available() else torch.device('cpu') dataset = BoxyDataset("/media/karsten/data_disk/boxy/boxy_images_raw", get_transforms(), constants.TRAIN_LABEL_FILE) dataset_valid = BoxyDataset( "/media/karsten/data_disk/boxy/boxy_images_raw", get_transforms(), constants.VALID_LABEL_FILE) data_loader = torch.utils.data.DataLoader(dataset, batch_size=4, shuffle=True, num_workers=4, collate_fn=utils.collate_fn) data_loader_valid = torch.utils.data.DataLoader( dataset_valid, batch_size=1, shuffle=False, num_workers=4, collate_fn=utils.collate_fn) model = faster_rcnn.fasterrcnn_resnet50_fpn(pretrained=True) num_classes = 2 # vehicle and background. Do I need background? in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = faster_rcnn.FastRCNNPredictor( in_features, num_classes) model.to(device) params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005) lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1) num_epochs = 10 for epoch in range(num_epochs): train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10) lr_scheduler.step() evaluate(model, data_loader_valid, device=device) torch.save(model, f"{epoch}_trained.pkl") print("Done training")
def test_faster_rcnn(self): model = faster_rcnn.fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=True, min_size=200, max_size=300) model.eval() x1 = torch.randn(3, 200, 300, requires_grad=True) x2 = torch.randn(3, 200, 300, requires_grad=True) self.run_test(model, ([x1, x2], ), rtol=1e-3, atol=1e-5) self.run_test( model, ([x1, x2], ), input_names=["images_tensors"], output_names=["outputs"], dynamic_axes={ "images_tensors": [0, 1, 2], "outputs": [0, 1, 2] }, rtol=1e-3, atol=1e-5, ) dummy_image = [torch.ones(3, 100, 100) * 0.3] images, test_images = _get_test_images() self.run_test( model, (images, ), additional_test_inputs=[(images, ), (test_images, ), (dummy_image, )], input_names=["images_tensors"], output_names=["outputs"], dynamic_axes={ "images_tensors": [0, 1, 2], "outputs": [0, 1, 2] }, rtol=1e-3, atol=1e-5, ) self.run_test( model, (dummy_image, ), additional_test_inputs=[(dummy_image, ), (images, )], input_names=["images_tensors"], output_names=["outputs"], dynamic_axes={ "images_tensors": [0, 1, 2], "outputs": [0, 1, 2] }, rtol=1e-3, atol=1e-5, )
def model(num_classes: int, backbone: Optional[nn.Module] = None, remove_internal_transforms: bool = True, pretrained: bool = True, **faster_rcnn_kwargs) -> nn.Module: """FasterRCNN model implemented by torchvision. # Arguments num_classes: Number of classes. backbone: Backbone model to use. Defaults to a resnet50_fpn model. remove_internal_transforms: The torchvision model internally applies transforms like resizing and normalization, but we already do this at the `Dataset` level, so it's safe to remove those internal transforms. pretrained: Argument passed to `fastercnn_resnet50_fpn` if `backbone is None`. By default it is set to True: this is generally used when training a new model (transfer learning). `pretrained = False` is used during inference (prediction) for cases where the users have their own pretrained weights. **faster_rcnn_kwargs: Keyword arguments that internally are going to be passed to `torchvision.models.detection.faster_rcnn.FastRCNN`. # Returns A Pytorch `nn.Module`. """ if backbone is None: model = fasterrcnn_resnet50_fpn(pretrained=pretrained, pretrained_backbone=pretrained, **faster_rcnn_kwargs) in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor( in_features, num_classes) backbone_param_groups = resnet_fpn.param_groups(model.backbone) else: model = FasterRCNN(backbone, num_classes=num_classes, **faster_rcnn_kwargs) backbone_param_groups = backbone.param_groups() patch_param_groups(model=model, backbone_param_groups=backbone_param_groups) if remove_internal_transforms: remove_internal_model_transforms(model) return model
def __init__( self, learning_rate: float = 0.0001, num_classes: int = 91, backbone: str = None, fpn: bool = True, pretrained_backbone: str = None, trainable_backbone_layers: int = 3, **kwargs, ): """ Args: learning_rate: the learning rate num_classes: number of detection classes (including background) pretrained: if true, returns a model pre-trained on COCO train2017 pretrained_backbone (str): if "imagenet", returns a model with backbone pre-trained on Imagenet trainable_backbone_layers: number of trainable resnet layers starting from final block """ super().__init__() self.learning_rate = learning_rate self.num_classes = num_classes self.backbone = backbone if backbone is None: self.model = fasterrcnn_resnet50_fpn( pretrained=True, trainable_backbone_layers=trainable_backbone_layers, ) in_features = self.model.roi_heads.box_predictor.cls_score.in_features self.model.roi_heads.box_predictor = FastRCNNPredictor( in_features, self.num_classes) else: backbone_model = create_fastercnn_backbone( self.backbone, fpn, pretrained_backbone, trainable_backbone_layers, **kwargs, ) self.model = FasterRCNN(backbone_model, num_classes=num_classes, **kwargs)
def faster_rcnn_res50(input_size, num_classes, self_pretrained): """ torchvision 自带的 fasterrcnn_resnet50_fpn 在 coco 上训练的,只提供了 resnet50 函数内部 resnet_fpn_backbone() 设置 backbone 类型 定义 layer2,3,4 之前层 require_grad=False, 之后 BackboneWithFPN() 添加了 FPN """ if not self_pretrained: print('load res50 backbone pretrained on COCO') model = fasterrcnn_resnet50_fpn( pretrained=not self_pretrained, min_size=input_size[0], max_size=input_size[1], # rpn_anchor_generator=anchor_generator, ) in_features = model.roi_heads.box_predictor.cls_score.in_features # 更新 `fasterrcnn_resnet50_fpn` 内 FastRCNNPredictor, num_classes 自定义 model.roi_heads.box_predictor = FastRCNNPredictor(in_channels=in_features, num_classes=num_classes) return model
def build_model(nclasses): """ Builds model. Uses Faster R-CNN pre-trained on COCO dataset. Args: nclasses: number of classes Return: model: Faster R-CNN pre-trained model """ # load pre-trained model on COCO model = fasterrcnn_resnet50_fpn(pretrained=False) # get the number of input features for the classifier in_features = model.roi_heads.box_predictor.cls_score.in_features # replace the pre-trained head with a new one model.roi_heads.box_predictor = FastRCNNPredictor(in_features, nclasses) return model
transforms.append(T.RandomHorizontalFlip(0.5)) return T.Compose(transforms) # === pathes pathes = {} pathes['project_root'] = '../' pathes['data'] = os.path.join(pathes['project_root'], 'materials/images/') # === load image img_path = os.path.join(pathes['data'], '2007_000720.jpg') img = Image.open(img_path) # === load model device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') model = fasterrcnn_resnet50_fpn(pretrained=True) model.to(device) # === transform image tr_imgs = get_transform(False) img_tensor, _ = tr_imgs(img, None) # === run model model.eval() with torch.no_grad(): prediction = model([img_tensor.to(device)]) predict_boxes = prediction[0]['boxes'].cpu().numpy() predict_labels = prediction[0]['labels'].cpu().numpy() predict_scores = prediction[0]['scores'].cpu().numpy()
def main(args): torch.cuda.set_device(0) random.seed(0) torch.manual_seed(0) torch.cuda.manual_seed(0) torch.cuda.manual_seed_all(0) utils.init_distributed_mode(args) print(args) device = torch.device(args.device) # Data loading code print("Loading data") if 'voc2007' in args.dataset: dataset, num_classes = get_dataset(args.dataset, "trainval", get_transform(train=True), args.data_path) dataset_test, _ = get_dataset(args.dataset, "test", get_transform(train=False), args.data_path) else: dataset, num_classes = get_dataset(args.dataset, "train", get_transform(train=True), args.data_path) dataset_test, _ = get_dataset(args.dataset, "val", get_transform(train=False), args.data_path) print("Creating data loaders") num_images = len(dataset) if 'voc' in args.dataset: init_num = 1000 budget_num = 1000 if 'retina' in args.model: init_num = 1000 budget_num = 500 else: init_num = 5000 budget_num = 1000 indices = list(range(num_images)) random.shuffle(indices) labeled_set = indices[:init_num] unlabeled_set = indices[init_num:] train_sampler = SubsetRandomSampler(labeled_set) test_sampler = torch.utils.data.SequentialSampler(dataset_test) data_loader_test = DataLoader(dataset_test, batch_size=1, sampler=test_sampler, num_workers=args.workers, collate_fn=utils.collate_fn) for cycle in range(args.cycles): if args.aspect_ratio_group_factor >= 0: group_ids = create_aspect_ratio_groups( dataset, k=args.aspect_ratio_group_factor) train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids, args.batch_size) else: train_batch_sampler = torch.utils.data.BatchSampler( train_sampler, args.batch_size, drop_last=True) data_loader = torch.utils.data.DataLoader( dataset, batch_sampler=train_batch_sampler, num_workers=args.workers, collate_fn=utils.collate_fn) print("Creating model") if 'voc' in args.dataset: if 'faster' in args.model: task_model = fasterrcnn_resnet50_fpn(num_classes=num_classes, min_size=600, max_size=1000) elif 'retina' in args.model: task_model = retinanet_resnet50_fpn(num_classes=num_classes, min_size=600, max_size=1000) else: if 'faster' in args.model: task_model = fasterrcnn_resnet50_fpn(num_classes=num_classes, min_size=800, max_size=1333) elif 'retina' in args.model: task_model = retinanet_resnet50_fpn(num_classes=num_classes, min_size=600, max_size=1000) task_model.to(device) if not args.init and cycle == 0 and args.skip: if 'faster' in args.model: checkpoint = torch.load(os.path.join( args.first_checkpoint_path, '{}_frcnn_1st.pth'.format(args.dataset)), map_location='cpu') elif 'retina' in args.model: checkpoint = torch.load(os.path.join( args.first_checkpoint_path, '{}_retinanet_1st.pth'.format(args.dataset)), map_location='cpu') task_model.load_state_dict(checkpoint['model']) # if 'coco' in args.dataset: # coco_evaluate(task_model, data_loader_test) # elif 'voc' in args.dataset: # voc_evaluate(task_model, data_loader_test, args.dataset) print("Getting stability") random.shuffle(unlabeled_set) if 'coco' in args.dataset: subset = unlabeled_set[:5000] else: subset = unlabeled_set # Update the labeled dataset and the unlabeled dataset, respectively labeled_set += subset[:budget_num] labeled_set = list(set(labeled_set)) # with open("vis/cycle_{}.txt".format(cycle), "rb") as fp: # Unpickling # labeled_set = pickle.load(fp) unlabeled_set = list(set(indices) - set(labeled_set)) # Create a new dataloader for the updated labeled dataset train_sampler = SubsetRandomSampler(labeled_set) continue params = [p for p in task_model.parameters() if p.requires_grad] task_optimizer = torch.optim.SGD(params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) task_lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( task_optimizer, milestones=args.lr_steps, gamma=args.lr_gamma) # Start active learning cycles training if args.test_only: if 'coco' in args.dataset: coco_evaluate(task_model, data_loader_test) elif 'voc' in args.dataset: voc_evaluate(task_model, data_loader_test, args.dataset) return print("Start training") start_time = time.time() for epoch in range(args.start_epoch, args.total_epochs): train_one_epoch(task_model, task_optimizer, data_loader, device, cycle, epoch, args.print_freq) task_lr_scheduler.step() # evaluate after pre-set epoch if (epoch + 1) == args.total_epochs: if 'coco' in args.dataset: coco_evaluate(task_model, data_loader_test) elif 'voc' in args.dataset: voc_evaluate(task_model, data_loader_test, args.dataset, path=args.results_path) if not args.skip and cycle == 0: if 'faster' in args.model: utils.save_on_master( { 'model': task_model.state_dict(), 'args': args }, os.path.join(args.first_checkpoint_path, '{}_frcnn_1st.pth'.format(args.dataset))) elif 'retina' in args.model: utils.save_on_master( { 'model': task_model.state_dict(), 'args': args }, os.path.join(args.first_checkpoint_path, '{}_retinanet_1st.pth'.format(args.dataset))) random.shuffle(unlabeled_set) # Update the labeled dataset and the unlabeled dataset, respectively labeled_set += unlabeled_set[:budget_num] labeled_set = list(set(labeled_set)) unlabeled_set = unlabeled_set[budget_num:] # Create a new dataloader for the updated labeled dataset train_sampler = SubsetRandomSampler(labeled_set) total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str))
def main(args): torch.cuda.set_device(0) random.seed(0) torch.manual_seed(0) torch.cuda.manual_seed(0) torch.cuda.manual_seed_all(0) print(args) device = torch.device(args.device) # Data loading code print("Loading data") if 'voc2007' in args.dataset: dataset, num_classes = get_dataset(args.dataset, "trainval", get_transform(train=True), args.data_path) dataset_test, _ = get_dataset(args.dataset, "test", get_transform(train=False), args.data_path) else: dataset, num_classes = get_dataset(args.dataset, "train", get_transform(train=True), args.data_path) dataset_test, _ = get_dataset(args.dataset, "val", get_transform(train=False), args.data_path) if 'voc' in args.dataset: init_num = 500 budget_num = 500 if 'retina' in args.model: init_num = 1000 budget_num = 500 else: init_num = 5000 budget_num = 1000 print("Creating data loaders") num_images = len(dataset) indices = list(range(num_images)) random.shuffle(indices) labeled_set = indices[:init_num] unlabeled_set = list(set(indices) - set(labeled_set)) train_sampler = SubsetRandomSampler(labeled_set) unlabeled_sampler = SubsetRandomSampler(unlabeled_set) test_sampler = torch.utils.data.SequentialSampler(dataset_test) data_loader_test = DataLoader(dataset_test, batch_size=1, sampler=test_sampler, num_workers=args.workers, collate_fn=utils.collate_fn) for cycle in range(args.cycles): if args.aspect_ratio_group_factor >= 0: group_ids = create_aspect_ratio_groups( dataset, k=args.aspect_ratio_group_factor) train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids, args.batch_size) unlabeled_batch_sampler = GroupedBatchSampler( unlabeled_sampler, group_ids, args.batch_size) else: train_batch_sampler = torch.utils.data.BatchSampler( train_sampler, args.batch_size, drop_last=True) unlabeled_batch_sampler = torch.utils.data.BatchSampler( unlabeled_sampler, args.batch_size, drop_last=True) data_loader = torch.utils.data.DataLoader( dataset, batch_sampler=train_batch_sampler, num_workers=args.workers, collate_fn=utils.collate_fn) unlabeled_dataloader = torch.utils.data.DataLoader( dataset, batch_sampler=unlabeled_batch_sampler, num_workers=args.workers, collate_fn=utils.collate_fn) print("Creating model") if 'voc' in args.dataset: if 'faster' in args.model: task_model = fasterrcnn_resnet50_fpn(num_classes=num_classes, min_size=600, max_size=1000) elif 'retina' in args.model: task_model = retinanet_resnet50_fpn(num_classes=num_classes, min_size=600, max_size=1000) else: if 'faster' in args.model: task_model = fasterrcnn_resnet50_fpn(num_classes=num_classes, min_size=800, max_size=1333) elif 'retina' in args.model: task_model = retinanet_resnet50_fpn(num_classes=num_classes, min_size=800, max_size=1333) task_model.to(device) params = [p for p in task_model.parameters() if p.requires_grad] task_optimizer = torch.optim.SGD(params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) task_lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( task_optimizer, milestones=args.lr_steps, gamma=args.lr_gamma) vae = VAE() params = [p for p in vae.parameters() if p.requires_grad] vae_optimizer = torch.optim.SGD(params, lr=args.lr / 10, momentum=args.momentum, weight_decay=args.weight_decay) vae_lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( vae_optimizer, milestones=args.lr_steps, gamma=args.lr_gamma) torch.nn.utils.clip_grad_value_(vae.parameters(), 1e5) vae.to(device) discriminator = Discriminator() params = [p for p in discriminator.parameters() if p.requires_grad] discriminator_optimizer = torch.optim.SGD( params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) discriminator_lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( discriminator_optimizer, milestones=args.lr_steps, gamma=args.lr_gamma) discriminator.to(device) # Start active learning cycles training if args.test_only: if 'coco' in args.dataset: coco_evaluate(task_model, data_loader_test) elif 'voc' in args.dataset: voc_evaluate(task_model, data_loader_test, args.dataset, False, path=args.results_path) return print("Start training") start_time = time.time() for epoch in range(args.start_epoch, args.total_epochs): train_one_epoch(task_model, task_optimizer, vae, vae_optimizer, discriminator, discriminator_optimizer, data_loader, unlabeled_dataloader, device, cycle, epoch, args.print_freq) task_lr_scheduler.step() vae_lr_scheduler.step() discriminator_lr_scheduler.step() # evaluate after pre-set epoch if (epoch + 1) == args.total_epochs: if 'coco' in args.dataset: coco_evaluate(task_model, data_loader_test) elif 'voc' in args.dataset: voc_evaluate(task_model, data_loader_test, args.dataset, False, path=args.results_path) # Update the labeled dataset and the unlabeled dataset, respectively random.shuffle(unlabeled_set) if 'coco' in args.dataset: subset = unlabeled_set[:10000] else: subset = unlabeled_set unlabeled_loader = DataLoader(dataset, batch_size=1, sampler=SubsetSequentialSampler(subset), num_workers=args.workers, pin_memory=True, collate_fn=utils.collate_fn) tobe_labeled_inds = sample_for_labeling(vae, discriminator, unlabeled_loader, budget_num) tobe_labeled_set = [subset[i] for i in tobe_labeled_inds] labeled_set += tobe_labeled_set unlabeled_set = list(set(unlabeled_set) - set(tobe_labeled_set)) # Create a new dataloader for the updated labeled dataset train_sampler = SubsetRandomSampler(labeled_set) unlabeled_sampler = SubsetRandomSampler(unlabeled_set) total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str))
def get_model(n_classes: int) -> FasterRCNN: model = fasterrcnn_resnet50_fpn(pretrained=False, progress=True, pretrained_backbone=True, num_classes=n_classes) return model
import torch import torchvision from torchvision.models.detection.faster_rcnn import fasterrcnn_resnet50_fpn import numpy as np from torchvision.ops import misc as misc_nn_ops # Regular resnet50, pretrained on ImageNet, without the classifier and the average pooling layer resnet50_1 = torch.nn.Sequential(*(list( torchvision.models.resnet50( pretrained=True, norm_layer=misc_nn_ops.FrozenBatchNorm2d).children()) [:-2])) resnet50_1.eval() # Resnet50, extract from the Faster R-CNN, also pre-trained on ImageNet resnet50_2 = fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=True).backbone.body resnet50_2.eval() # Loading a random image, converted to torch.Tensor, rescalled to [0, 1] (not that it matters) image = torch.ones((1, 3, 224, 224)) # Obtaining the model outputs with torch.no_grad(): # Output from the regular resnet50 output_1 = resnet50_1(image) # Output from the resnet50 extracted from the Faster R-CNN output_2 = resnet50_2(image)["3"] # Their outputs aren't the same, which I would assume they should be np.testing.assert_almost_equal(output_1.numpy(), output_2.numpy())