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())
