def get_model(classes=11):
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
backbone.out_channels = 1280
anchor_generator = AnchorGenerator(sizes=((16, 32, 64, 128, 256),),
aspect_ratios=((0.5, 1.0, 1.5),))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(
featmap_names=['0', '1', '2', '3'],
output_size=7,
sampling_ratio=2)
# put the pieces together inside a FasterRCNN model
model = FasterRCNN(backbone,
num_classes=classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
return model
def model_creation(pretrain=True,
num_classes=5,
num_epoch=20,
device="cuda:0"):
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(
pretrained=True)
num_classes = 5
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
backbone.out_channels = 1280
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
output_size=7,
sampling_ratio=2)
model = FasterRCNN(backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
if device == 'cuda:0':
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if device == "cpu":
print("cuda is no available")
model.to(device)
# construct an optimizer
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)
# and a learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.1)
# let's train it for 10 epochs
return model, optimizer, lr_scheduler, num_epoch,
def get_faster_RCNN(model_name, pretrained=True, num_classes=14, img_size=512):
backbone = Timm_model_featuremap(model_name,
img_size,
pretrained=pretrained)
anchor_generator = AnchorGenerator(sizes=((8, 8, 16, 16, 32, 32, 64, 64,
128, 128, 256, 512), ),
aspect_ratios=((0.2, 0.5, 0.7, 1.0, 1.5,
2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
output_size=7,
sampling_ratio=2)
model = FasterRCNN(backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
return model
def fasterrcnn_mobilenetv2(num_classes, **kwargs):
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
backbone.out_channels = 1280
anchor_generator = AnchorGenerator(
sizes=((32, 64, 128, 256, 512),), aspect_ratios=((0.5, 1.0, 2.0),)
)
roi_pooler = torchvision.ops.MultiScaleRoIAlign(
featmap_names=["0"], output_size=7, sampling_ratio=2
)
model = FasterRCNN(
backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler,
**kwargs
)
return model
def fasterrcnn_resnet101_fpn(pretrained=False, progress=False,
num_classes=91, pretrained_backbone=True, **kwargs):
"""
Constructs a Faster R-CNN model with a ResNet-101-FPN backbone.
The input to the model is expected to be a list of tensors, each of shape ``[C, H, W]``, one for each
image, and should be in ``0-1`` range. Different images can have different sizes.
The behavior of the model changes depending if it is in training or evaluation mode.
During training, the model expects both the input tensors, as well as a targets (list of dictionary),
containing:
- boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with values
between ``0`` and ``H`` and ``0`` and ``W``
- labels (``Int64Tensor[N]``): the class label for each ground-truth box
The model returns a ``Dict[Tensor]`` during training, containing the classification and regression
losses for both the RPN and the R-CNN.
During inference, the model requires only the input tensors, and returns the post-processed
predictions as a ``List[Dict[Tensor]]``, one for each input image. The fields of the ``Dict`` are as
follows:
- boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with values between
``0`` and ``H`` and ``0`` and ``W``
- labels (``Int64Tensor[N]``): the predicted labels for each image
- scores (``Tensor[N]``): the scores or each prediction
Example::
>>> model = torchvision.models.detection.fasterrcnn_resnet101_fpn(pretrained=False)
>>> model.eval()
>>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
>>> predictions = model(x)
Arguments:
pretrained (bool): If True, returns a model pre-trained on COCO train2017(currently don't have pre-trained models)
progress (bool): If True, displays a progress bar of the download to stderr
"""
if pretrained:
raise Exception("resnet101 cannot use pre-trained models")
# no need to download the backbone if pretrained is set
pretrained_backbone = False
backbone = resnet_fpn_backbone('resnet101', pretrained_backbone)
model = FasterRCNN(backbone, num_classes, **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls['fasterrcnn_resnet101_fpn_coco'],
progress=progress)
model.load_state_dict(state_dict)
return model
def get_model():
# model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
# in_features = model.roi_heads.box_predictor.cls_score.in_features
# model.roi_heads.box_predictor = FastRCNNPredictor(in_features, NUM_CLASSES)
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
backbone.out_channels = 1280
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
model = FasterRCNN(backbone,
num_classes=NUM_CLASSES,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
return model
def __init__(self):
super().__init__()
backbone = torchvision.models.vgg16(pretrained=False).features
backbone.out_channels = 512
anchor_sizes = ((8, 16, 32, 64, 128, 256, 512), )
aspect_ratios = ((1 / 2, 1 / 3, 1 / 4, 1 / 5, 1 / 6, 1 / math.sqrt(2),
1, 2, math.sqrt(2), 3, 4, 5, 6, 7, 8), )
anchor_generator = AnchorGenerator(sizes=anchor_sizes,
aspect_ratios=aspect_ratios)
roi_pooler = torchvision.ops.MultiScaleRoIAlign(
featmap_names=['0', '1', '2', '3', '4'],
output_size=7,
sampling_ratio=2)
self.model = FasterRCNN(backbone,
num_classes=7,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
self.device = torch.device('cpu')
self.model.load_state_dict(torch.load('2.pth'))
self.model.to(self.device)
self.model.eval()
def __init__(self, out_dim=None, fine_tune=False):
super(FasterRCNN_Encoder, self).__init__()
backbone = resnet_fpn_backbone('resnet50', False)
self.faster_rcnn = FasterRCNN(backbone,
num_classes=91,
rpn_post_nms_top_n_train=200,
rpn_post_nms_top_n_test=100)
state_dict = load_state_dict_from_url(
model_urls['fasterrcnn_resnet50_fpn_coco'], progress=True)
self.faster_rcnn.load_state_dict(state_dict)
# modify the last linear layer of the ROI pooling if there is
# a special requirement of output size
if out_dim is not None:
self.faster_rcnn.roi_heads.box_head.fc7 = nn.Linear(
in_features=1024, out_features=out_dim)
# in captioning task, we may not want fine-tune faster-rcnn model
if not fine_tune:
for param in self.faster_rcnn.parameters():
param.requires_grad = False
def official_faster(num_classes):
import torchvision
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
# load a pre-trained model for classification and return
# only the features
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
# FasterRCNN needs to know the number of
# output channels in a backbone. For mobilenet_v2, it's 1280
# so we need to add it here
backbone.out_channels = 1280
# here mobilenet_v2 just has features, and classifier, has not layer1-layer4, so can'y add fpn
# let's make the RPN generate 5 x 3 anchors per spatial
# location, with 5 different sizes and 3 different aspect
# ratios. We have a Tuple[Tuple[int]] because each feature
# map could potentially have different sizes and
# aspect ratios
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),
aspect_ratios=((0.5, 1.0, 2.0),))
# let's define what are the feature maps that we will
# use to perform the region of interest cropping, as well as
# the size of the crop after rescaling.
# if your backbone returns a Tensor, featmap_names is expected to
# be [0]. More generally, the backbone should return an
# OrderedDict[Tensor], and in featmap_names you can choose which
# feature maps to use.
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
# put the pieces together inside a FasterRCNN model
model = FasterRCNN(backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
return model
def __init__(self, load_weights=False):
super(DecideNet, self).__init__()
self.seen = 0
self.reg_net_feat = [(7, 20, 1), (5, 40, 1), (5, 20, 1), (5, 10, 1),
(1, 1, 1)]
self.reg_net = make_layers(self.reg_net_feat, 3)
self.quality_net_feat = [(7, 20, 1), (5, 40, 1), (5, 20, 1), (1, 1, 1)]
self.quality_net = make_layers(self.quality_net_feat, 3)
# backbone = torchvision.models.mobilenet_v2(pretrained=False).features
backbone = torchvision.models.detection.backbone_utils.resnet_fpn_backbone(
'resnet101', 'False')
# model = torchvision.models.detection.fasterrcnn_resnet50_fpn()
# backbone.out_channels = 1280
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
self.fastercnn = FasterRCNN(backbone,
num_classes=2,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
def other():
# 在COCO上加载经过预训练的预训练模型
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(
pretrained=True)
# replace the classifier with a new one, that has
# 将分类器替换为具有用户定义的 num_classes的新分类器
num_classes = 2 # 1 class (person) + background
# 获取分类器的输入参数的数量
in_features = model.roi_heads.box_predictor.cls_score.in_features
# 用新的头部替换预先训练好的头部
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
# 加载预先训练的模型进行分类和返回
# 只有功能
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
# FasterRCNN需要知道骨干网中的输出通道数量。对于mobilenet_v2，它是1280，所以我们需要在这里添加它
backbone.out_channels = 1280
# 我们让RPN在每个空间位置生成5 x 3个锚点
# 具有5种不同的大小和3种不同的宽高比。
# 我们有一个元组[元组[int]]
# 因为每个特征映射可能具有不同的大小和宽高比
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
# 定义一下我们将用于执行感兴趣区域裁剪的特征映射，以及重新缩放后裁剪的大小。
# 如果您的主干返回Tensor，则featmap_names应为[0]。
# 更一般地，主干应该返回OrderedDict [Tensor]
# 并且在featmap_names中，您可以选择要使用的功能映射。
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
# 将这些pieces放在FasterRCNN模型中
model = FasterRCNN(backbone,
num_classes=2,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
def get_model_instance_segmentation6(num_classes):
backbone = torchvision.models.squeezenet1_1(pretrained=False).features
backbone.out_channels = 512
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=1,
sampling_ratio=2)
model = FasterRCNN(backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
print(
"get_model_instance_segmentation6 call6 - out_channels :512, 4,808,441 / (15,000,000) "
)
return model
def get_faster_rcnn_101_model_instance(num_classes=15, min_size=512):
# https://discuss.pytorch.org/t/faster-mask-rcnn-rpn-custom-anchorgenerator/69962/2
sizes = ((2, ), (4, ), (8, ), (16, ), (32, ), (64, ), (128, ), (256, ),
(512, ))
aspect_ratios = tuple([(0.5, 1.0, 2.0) for _ in range(len(sizes))])
anchor_generator = AnchorGenerator(sizes=sizes,
aspect_ratios=aspect_ratios)
# define backbone model
backbone = resnet_fpn_backbone('resnet101', True)
model = FasterRCNN(backbone,
num_classes,
min_size=min_size,
rpn_anchor_generator=anchor_generator)
# 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, num_classes)
return model
def __init__(self, hparams):
super().__init__()
self.hparams = hparams
self.output_dim = 800 * 800
#self.kernel_size = 4
# TODO: add pretrained weight path
# TODO: remove this to train models again
d = dict(latent_dim=64, hidden_dim=128, batch_size=16)
hparams2 = Namespace(**d)
# ------------------
# PRE-TRAINED MODEL
# ------------------
ae = BasicAE.load_from_checkpoint(self.hparams.pretrained_path)
# ae = BasicAE(hparams2)
ae.freeze()
self.backbone = ae.encoder
self.backbone.c3_only = True
self.backbone.out_channels = 32
# ------------------
# FAST RCNN
# ------------------
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
output_size=7,
sampling_ratio=2)
self.fast_rcnn = FasterRCNN(self.backbone,
num_classes=9,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
# for unfreezing encoder later
self.frozen = True
def get_pretrain_model(model_name, num_classes):
if model_name == 'faster_rcnn':
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(
pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(
in_features, num_classes)
elif model_name == 'mask_rcnn':
model = torchvision.models.detection.maskrcnn_resnet50_fpn(
pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(
in_features, num_classes)
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
model.roi_heads.mask_predictor = MaskRCNNPredictor(
in_features_mask, hidden_layer, num_classes)
else:
backbone = torchvision.models.mobilenet_v2().features
backbone.out_channels = 1280
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
model = FasterRCNN(backbone,
num_classes=num_classes,
box_roi_pool=roi_pooler,
rpn_anchor_generator=anchor_generator)
return model
def get_model(pretrained=True):
# # load a model pre-trained on COCO
# model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=pretrain, min_size=540, max_size=960)
# # replace the classifier with a new one, that has
# # num_classes which is user-defined
# num_classes = 13
# # get number of input features for the classifier
# in_features = model.roi_heads.box_predictor.cls_score.in_features
# # replace the pre-train head with a new one
# model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
backbone = resnet_fpn_backbone("resnet50", pretrained)
# backbone.out_channels = 256
# anchor_size =((8,), (16,), (32,), (64,), (128,), (256,), (512,))
anchor_size = ((16,), (32,), (64,), (128,), (256,))
aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_size)
rpn_anchor_generator = AnchorGenerator(anchor_size, aspect_ratios)
model = FasterRCNN(backbone=backbone,
num_classes=13,
rpn_anchor_generator=rpn_anchor_generator,
min_size=540,
max_size=900)
return model
def set_basemodel_in_framework(basemodel,
framework,
num_classes,
rpn_anchor_generator=None,
box_roi_pooler=None
):
"""
FasterRCNN frameworks require rpn_anchor or box_roi_pool to be
specified. If they're not specified PyTorch will just use defaults which
are fine.
If a custom rpn_anchor_generator/roi_pooler is desired, define them as
follows (example)...
anchor_generator = torchvision.models.detection.rpn.AnchorGenerator(
sizes=((32, 64, 128, 256, 512),),
aspect_ratios=((0.5, 1.0, 2.0),))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=["0"],
output_size=7,
sampling_ratio=2)
... and pass them into the function
SSD models don't need rpn anchor generators or roi poolers
"""
if framework == "FasterRCNN":
model = FasterRCNN(basemodel,
num_classes=num_classes,
rpn_anchor_generator=None,
box_roi_pool=None)
elif framework == "SSD":
model = SSD(basemodel, num_classes=num_classes)
return model
def get_model_instance_segmentation5(num_classes):
# COCO 에서 미리 학습된 인스턴스 분할 모델을 읽어옵니다
#model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False)
#model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False)
backbone = torchvision.models.densenet161(pretrained=False).features
#backbone.out_channels = 1
backbone.out_channels = 256
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=1,
sampling_ratio=2)
model = FasterRCNN(backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
print("densenet161 call2 - out_channels :256, 28,506,873 / 150M")
# 분류를 위한 입력 특징 차원을 얻습니다
#in_features = backbone
# 미리 학습된 헤더를 새로운 것으로 바꿉니다
#model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
#in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
#hidden_layer = 1
# and replace the mask predictor with a new one
#model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
# hidden_layer,
# num_classes)
return model
def set_model(self):
"""
Set model and determine configuration
:return: None, generate self.model to be used for training and testing
"""
# Default values: box_score_thresh = 0.05, box_nms_thresh = 0.5
kwargs = {
'box_score_thresh': 0.3,
'box_nms_thresh': 0.3,
'box_detections_per_img': 6
}
# self.model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False,
# pretrained_backbone=True,
# **kwargs)
self.model = FasterRCNN(self.backbone, num_classes=7, **kwargs)
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
num_classes = 7
in_features = self.model.roi_heads.box_predictor.cls_score.in_features
self.model.roi_heads.box_predictor = FastRCNNPredictor(
in_features, num_classes)
# Allow Multiple GPUs:
# if torch.cuda.device_count() > 1:
# self.model = nn.DataParallel(self.model)
self.model = self.model.to(device)
if self.params is None:
params = [p for p in self.model.parameters() if p.requires_grad]
else:
# TODO: Enable user defined model params
pass
self.optimizer = torch.optim.SGD(params, lr=0.01)
x = self.relu(x)
x = self.bn1(x)
x = self.maxpoolH_0(x)
x1 = self.the_block1(x)
x2 = self.the_block2(x1)
x3 = self.the_block3(x2)
x4 = self.the_block4(x3)
# print("target ", annotations)
# print("target size : ", annotations.shape)
# print("x4 size : ", x4.shape)
# print("="*50)
return x4
if __name__ == "__main__":
backbone = FractalNet(11, BasicBlock, BigBlock, True)
images, boxes = torch.rand(1, 3, 800, 1200), torch.rand(1, 11, 4)
labels = torch.randint(1, 91, (1, 11))
# images = list(image for image in images)
targets = []
for i in range(len(images)):
d = {}
d['boxes'] = boxes[i]
d['labels'] = labels[i]
targets.append(d)
# output = backbone((images, targets))
model = FasterRCNN(backbone, num_classes=11)
output = model((images, targets))
"normal/part2/"
]
ears = ["right ear/", "left ear/"]
transform_img = transforms.Compose([transforms.ToTensor()])
EPOCH = 250
CLASSES = 3
DEVICE = torch.device("cuda")
BATCH_SIZE = 10
anchor_generator = AnchorGenerator(sizes=((32, 64), ),
aspect_ratios=((0.6, 1.0, 1.6), ))
backbone = torchvision.models.vgg19(pretrained=False).features
backbone.out_channels = 512
model = FasterRCNN(backbone,
num_classes=CLASSES,
rpn_anchor_generator=anchor_generator)
model.load_state_dict(
torch.load('models_new/' + 'model_' + str(EPOCH) + '.pth'))
model.to(DEVICE)
model.eval()
start_time = time.time()
ear_count = 0
for T in types:
for E in ears:
CTs = os.listdir(data_path + dataset_name + T + E)
for CT in CTs:
print('current path:{}'.format(data_path + dataset_name + T + E +
CT))
ear_count += 1
img_names = glob.glob(data_path + dataset_name + T + E + CT +
def main(network):
# train on the GPU or on the CPU, if a GPU is not available
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
# our dataset has two classes only - background and person
num_classes = 2
#dataset = torch.utils.data.Subset(TBdata,[range(len(TBdata))])
indices = torch.randperm(len(TBdata)).tolist()
dataset = torch.utils.data.Subset(TBdata, indices[:])
indices_ = torch.randperm(len(TBdata_test)).tolist()
dataset_val = torch.utils.data.Subset(TBdata_test, indices_[:])
# get the model using our helper function
#model = get_model_instance_segmentation(num_classes)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=8,
sampler=None,
num_workers=0,
collate_fn=collate_fn)
dataloader_val = torch.utils.data.DataLoader(dataset_val,
batch_size=8,
sampler=None,
num_workers=0,
collate_fn=collate_fn)
#Calculated statistics on training data:
#Transform parameters
min_size = 550
max_size = 700
image_means = [0.9492, 0.9492, 0.9492]
image_stds = [0.1158, 0.1158, 0.1158]
if network == 'resnet50':
backbone = resnet_fpn_backbone('resnet50', True)
model = FasterRCNN(backbone,
num_classes,
min_size=min_size,
max_size=max_size,
image_mean=image_means,
image_std=image_stds)
elif network == 'resnet18':
backbone = resnet_fpn_backbone('resnet18', True)
model = FasterRCNN(backbone,
num_classes,
min_size=min_size,
max_size=max_size,
image_mean=image_means,
image_std=image_stds)
elif network == 'resnet152':
backbone = resnet_fpn_backbone('resnet152', True)
model = FasterRCNN(backbone,
num_classes,
min_size=min_size,
max_size=max_size,
image_mean=image_means,
image_std=image_stds)
elif network == 'RPNresnet50':
backbone = resnet_fpn_backbone('resnet50', True)
model = RPN_custom(backbone,
num_classes,
min_size=min_size,
max_size=max_size,
image_mean=image_means,
image_std=image_stds)
elif network == 'RPNresnet152':
backbone = resnet_fpn_backbone('resnet152', True)
model = RPN_custom(backbone,
num_classes,
min_size=min_size,
max_size=max_size,
image_mean=image_means,
image_std=image_stds)
# move model to the right device
model.to(device)
# construct an optimizer
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)
# and a learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.1)
num_epochs = 10
Ls = {
'total loss': [],
'loss_classifier': [],
'loss_box_reg': [],
'loss_objectness': [],
'loss_rpn_box_reg': []
}
Ls_val = {
'total loss': [],
'loss_classifier': [],
'loss_box_reg': [],
'loss_objectness': [],
'loss_rpn_box_reg': []
}
for epoch in range(num_epochs):
# train for one epoch, printing every 10 iterations
train_one_epoch(model,
optimizer,
dataloader,
device,
epoch,
print_freq=10)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
#evaluate(model, dataloader_test, device=device)
Ls_val = record_losses(model, dataloader_val, device, Ls_val, network)
#record losses
Ls = record_losses(model, dataloader, device, Ls, network)
#If folder does not exist already, create it
output_loc = f'./{network}/'
if not os.path.exists(output_loc):
os.makedirs(output_loc)
torch.save(model.state_dict(), output_loc + 'model.pt')
print("That's it!")
return Ls, Ls_val, num_epochs
aspect_ratios=((0.5, 1.0, 2.0),))
# let's define what are the feature maps that we will
# use to perform the region of interest cropping, as well as
# the size of the crop after rescaling.
# if your backbone returns a Tensor, featmap_names is expected to
# be [0]. More generally, the backbone should return an
# OrderedDict[Tensor], and in featmap_names you can choose which
# feature maps to use.
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
# put the pieces together inside a FasterRCNN model
model = FasterRCNN(backbone,
num_classes=2,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
def get_model_instance_segmentation(num_classes):
# load an instance segmentation model pre-trained pre-trained on COCO
model = torchvision.models.detection.maskrcnn_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, num_classes)
# Initialize Backbone Res-34
sim_clr = models.simclr_resnet('res34', non_linear_head=False)
# Load weights for backbone from PIRL pretrained network.
sim_dict = args.backbone_pretrained_weights_file
#"../e1_simclr_auto_main_epoch_90"
sim_clr.load_state_dict(torch.load(sim_dict, map_location=device))
model_res34 = torchvision.models.resnet34(pretrained=False)
network_helpers.copy_weights_between_models(sim_clr, model_res34)
modules = list(model_res34.children())[:-1]
backbone = nn.Sequential(*modules)
backbone.out_channels = 512
model_fnn = FasterRCNN(backbone=backbone, num_classes=10)
# ImageNet
#in_features = model_fnn.roi_heads.box_predictor.cls_score.in_features
#model_fnn.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes=10)
model_fnn.to(device)
params_fnn = [p for p in model_fnn.parameters() if p.requires_grad]
params_cnn_fnn = list(params_cnn) + list(params_fnn)
optimizer = optim.SGD(params_cnn_fnn,
lr=args.lr,
momentum=0.6,
weight_decay=0.0005)
def get_fasterrcnn_model_swin(arch_str,
num_classes,
pretrained=False,
pretrained_backbone=True,
**kwargs):
"""Creates FasterRCNN model with swin transformer backbone"""
anchor_sizes = (
(32, ),
(64, ),
(128, ),
(256, ),
)
aspect_ratios = ((0.5, 1.0, 2.0), ) * len(anchor_sizes)
anchor_generator = AnchorGenerator(sizes=anchor_sizes,
aspect_ratios=aspect_ratios)
#roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0','1','2','3'],
# output_size=7,
# sampling_ratio=2)
img_size = 224 if arch_str in "swin_tiny swin_small".split() else 384
window_size = 7 if arch_str in "swin_tiny swin_small".split() else 12
depths = [2, 2, 6, 2] if arch_str == "swin_tiny" else [2, 2, 18, 2]
scale_factors = {
"swin_tiny": 1.0,
"swin_small": 1.5,
"swin_base": 2.0,
"swin_large": 2.0
}
sf = scale_factors[arch_str]
embed_dim = int(96 * sf)
fpn_cin = [int(96 * sf * 2**i) for i in range(4)]
#fpn_cin = [int(i*sf) for i in [96, 192, 384, 768]]
backbone = SwinTransformerFPN(img_size=img_size,
window_size=window_size,
embed_dim=embed_dim,
depths=depths,
fpn_cin=fpn_cin,
fpn_cout=256)
if pretrained_backbone:
sd = load_state_dict_from_url(_model_urls[f'{arch_str}_{img_size}'],
progress=True,
map_location=default_device())['model']
sd_model = backbone.state_dict()
sd = {k: v for k, v in sd.items() if k in sd_model.keys()}
sd_model.update(sd)
backbone.load_state_dict(sd_model)
model = FasterRCNN(
backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
#box_roi_pool=roi_pooler,
box_fg_iou_thresh=0.5,
box_bg_iou_thresh=0.5,
image_mean=[0.0, 0.0, 0.0], # already normalized by fastai
image_std=[1.0, 1.0, 1.0],
#min_size=IMG_SIZE,
#max_size=IMG_SIZE,
**kwargs)
return model.train()
def fasterrcnn_resnet50_fpn(pretrained=False,
progress=True,
num_classes=91,
pretrained_backbone=True,
**kwargs):
"""
Constructs a Faster R-CNN model with a ResNet-50-FPN backbone.
The input to the model is expected to be a list of tensors, each of shape ``[C, H, W]``, one for each
image, and should be in ``0-1`` range. Different images can have different sizes.
The behavior of the model changes depending if it is in training or evaluation mode.
During training, the model expects both the input tensors, as well as a targets (list of dictionary),
containing:
- boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with values of ``x``
between ``0`` and ``W`` and values of ``y`` between ``0`` and ``H``
- labels (``Int64Tensor[N]``): the class label for each ground-truth box
The model returns a ``Dict[Tensor]`` during training, containing the classification and regression
losses for both the RPN and the R-CNN.
During inference, the model requires only the input tensors, and returns the post-processed
predictions as a ``List[Dict[Tensor]]``, one for each input image. The fields of the ``Dict`` are as
follows:
- boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with values of ``x``
between ``0`` and ``W`` and values of ``y`` between ``0`` and ``H``
- labels (``Int64Tensor[N]``): the predicted labels for each image
- scores (``Tensor[N]``): the scores or each prediction
Faster R-CNN is exportable to ONNX for a fixed batch size with inputs images of fixed size.
Example::
>>> model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
>>> # For training
>>> images, boxes = torch.rand(4, 3, 600, 1200), torch.rand(4, 11, 4)
>>> labels = torch.randint(1, 91, (4, 11))
>>> images = list(image for image in images)
>>> targets = []
>>> for i in range(len(images)):
>>> d = {}
>>> d['boxes'] = boxes[i]
>>> d['labels'] = labels[i]
>>> targets.append(d)
>>> output = model(images, targets)
>>> # For inference
>>> model.eval()
>>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
>>> predictions = model(x)
>>>
>>> # optionally, if you want to export the model to ONNX:
>>> torch.onnx.export(model, x, "faster_rcnn.onnx", opset_version = 11)
Arguments:
pretrained (bool): If True, returns a model pre-trained on COCO train2017
progress (bool): If True, displays a progress bar of the download to stderr
"""
if pretrained:
# no need to download the backbone if pretrained is set
pretrained_backbone = False
backbone = resnet_fpn_backbone('resnet50', pretrained_backbone)
model = FasterRCNN(backbone, num_classes, **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(
model_urls['fasterrcnn_resnet50_fpn_coco'], progress=progress)
model.load_state_dict(state_dict)
return model
collate_fn=collate_fn)
valid_loader = DataLoader(valid_dataset,
batch_size=1,
shuffle=False,
num_workers=cfg.workers,
collate_fn=collate_fn)
# Start to traning FASTER RCNN
backbone = backboneNet_efficient() # use efficientnet as our backbone
backboneFPN = backboneWithFPN(backbone) # add FPN
anchor_generator = AnchorGenerator(cfg.anchor_sizes, cfg.aspect_ratios)
model_ft = FasterRCNN(backboneFPN,
num_classes=cfg.num_classes,
rpn_anchor_generator=anchor_generator,
min_size=cfg.min_size,
max_size=cfg.max_size)
model_ft.to(device)
optimizer_ft = optim.SGD(model_ft.parameters(),
lr=cfg.learning_rate,
momentum=cfg.momentum,
weight_decay=cfg.weight_decay)
lr_scheduler = lr_scheduler.MultiStepLR(optimizer_ft,
milestones=cfg.milestones,
gamma=cfg.gamma)
model_ft = train_model(model_ft,
def main():
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(
pretrained=True)
num_classes = 2
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
backbone.out_channels = 1280
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
model = FasterRCNN(backbone,
num_classes=2,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
# Train classifier
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
dataset = OysterDataset(
"/Users/darwin/Downloads/Bay project/Training",
"/Users/darwin/Downloads/Bay project/Training/TrainingDetectionLabels2.csv",
get_transform(train=True))
dataset_test = OysterDataset(
"/Users/darwin/Downloads/Bay project/Training",
"/Users/darwin/Downloads/Bay project/Training/TrainingDetectionLabels2.csv",
get_transform(train=False))
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:100])
dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=2,
shuffle=False,
num_workers=4,
collate_fn=utils.collate_fn)
data_loader_test = torch.utils.data.DataLoader(dataset_test,
batch_size=2,
shuffle=False,
num_workers=4,
collate_fn=utils.collate_fn)
model = get_model_instance_segmentation(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.8,
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_test, device=device)
print("That's it!")
def get_fasterrcnn_model(arch_str,
num_classes,
pretrained=True,
pretrained_backbone=True,
trainable_layers=5,
**kwargs):
"""Creates FasterRCNN model with resnet backbone"""
#if pretrained == True: pretrained_backbone=False
backbone = resnet_fpn_backbone(arch_str,
pretrained=pretrained_backbone,
trainable_layers=trainable_layers)
anchor_sizes = (
(16, ),
(32, ),
(64, ),
(128, ),
(256, ),
)
aspect_ratios = ((0.5, 1.0, 2.0), ) * len(anchor_sizes)
anchor_generator = AnchorGenerator(sizes=anchor_sizes,
aspect_ratios=aspect_ratios)
model = FasterRCNN(
backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_fg_iou_thresh=0.5,
box_bg_iou_thresh=0.5,
image_mean=[0.0, 0.0, 0.0], # already normalized by fastai
image_std=[1.0, 1.0, 1.0],
#min_size = 1,
#box_score_thresh=0.6,
**kwargs)
if pretrained:
try:
pretrained_dict = load_state_dict_from_url(
_model_urls['fasterrcnn_' + arch_str + '_fpn_coco'],
progress=True)
model_dict = model.state_dict()
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if (k in model_dict) and (
model_dict[k].shape == pretrained_dict[k].shape)
}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
#overwrite_eps(model, 0.0)
for module in model.modules():
if isinstance(module, FrozenBatchNorm2d):
module.eps = 0.0
except Exception as e:
#print(e)
print("No pretrained coco model found for fasterrcnn_" + arch_str)
print("This does not affect the backbone.")
return model.train()
def fasterrcnn_model(backbone,
class_num=2,
pool_layers_num=4,
pooled_size=7,
pretrained=True):
"""
pool_layers_num: MultiScaleRoIAlignで使う層の数. 'resnet50_coco'では無視される. 安全そうな範囲で1~4で指定
pooled_size: RoIPool後のmap size. 'resnet50_coco'では無視される. 安全そうな5~9で指定
"""
backbone_list = {
'resnet18': True,
'resnet34': True,
'resnet50': True,
'resnet50_coco': True, # いままでのやつ、headまでCOCOでpretrained
'resnet101': True, # batch_size=4は乗る
'resnet152': True, # batch_size=4は乗る
'resnext50_32x4d': True,
# 'resnext101_32x8d': True, # エラー起きる
# 'wide_resnet50_2': True, # エラー起きる
# 'wide_resnet101_2': True # エラー起きる
}
assert backbone in backbone_list.keys(
), 'Backbone\'s name is not valid. Available backbones: %s' % str(
list(backbone_list.keys()))
if backbone == 'resnet50_coco':
# 今まで使っていたmodel、headまでpretrainedでweightsを読み込んでおり構造は弄れない
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(
pretrained=pretrained)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(
in_features, class_num)
else:
# backboneだけpretrained
assert pool_layers_num in [
1, 2, 3, 4
], 'pool_layers_num must be in [1, 2, 3, 4] You selected %d' % (
pool_layers_num)
assert pooled_size in [
5, 6, 7, 8, 9
], 'pooled_size must be in [5, 6, 7, 8, 9] You selected %d' % (
pooled_size)
# anchor_sizesはデフォルトから1スケール落とした。 default: ((32,), (64,), (128,), (256,), (512,))
anchor_sizes = ((16), (32, ), (64, ), (128, ), (256, ))
# anchor_ratiosは4:1の比を追加
aspect_ratios = ((0.25, 0.5, 1.0, 2.0, 4.0), ) * len(anchor_sizes)
anchor_generator = AnchorGenerator(sizes=anchor_sizes,
aspect_ratios=aspect_ratios)
# デフォルトでマルチスケールのRoIAlignになっている。headに近い4層から特徴を抽出しているはず
roi_pooler = torchvision.ops.MultiScaleRoIAlign(
featmap_names=[str(n) for n in range(pool_layers_num)],
output_size=pooled_size,
sampling_ratio=2)
backbone = resnet_fpn_backbone(backbone, pretrained=pretrained)
model = FasterRCNN(backbone,
num_classes=class_num,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
return model