import torch
import torchvision
from models.deeplabv3plus import DeepLabV3plus
from models.resnet import build_resnet
encoder = build_resnet('18', output_stride=8)
model = DeepLabV3plus(encoder, 12)
from datasets.segmentation import CamVid
dataset = CamVid("../seg_datasets/CamVid")
image1 = torchvision.transforms.functional.to_tensor(dataset[0][0])
image2 = torchvision.transforms.functional.to_tensor(dataset[1][0])
images = torch.stack([image1, image2], dim=0)
print('Without gradient checkpointing')
model(images, False)
print('\nWith gradient checkpointing')
model(images, True)
def __init__(self, cfg):
super(SeqNet, self).__init__()
backbone, box_head = build_resnet(name="resnet50", pretrained=True)
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
head = RPNHead(
in_channels=backbone.out_channels,
num_anchors=anchor_generator.num_anchors_per_location()[0],
)
pre_nms_top_n = dict(training=cfg.MODEL.RPN.PRE_NMS_TOPN_TRAIN,
testing=cfg.MODEL.RPN.PRE_NMS_TOPN_TEST)
post_nms_top_n = dict(training=cfg.MODEL.RPN.POST_NMS_TOPN_TRAIN,
testing=cfg.MODEL.RPN.POST_NMS_TOPN_TEST)
rpn = RegionProposalNetwork(
anchor_generator=anchor_generator,
head=head,
fg_iou_thresh=cfg.MODEL.RPN.POS_THRESH_TRAIN,
bg_iou_thresh=cfg.MODEL.RPN.NEG_THRESH_TRAIN,
batch_size_per_image=cfg.MODEL.RPN.BATCH_SIZE_TRAIN,
positive_fraction=cfg.MODEL.RPN.POS_FRAC_TRAIN,
pre_nms_top_n=pre_nms_top_n,
post_nms_top_n=post_nms_top_n,
nms_thresh=cfg.MODEL.RPN.NMS_THRESH,
)
faster_rcnn_predictor = FastRCNNPredictor(2048, 2)
reid_head = deepcopy(box_head)
box_roi_pool = MultiScaleRoIAlign(featmap_names=["feat_res4"],
output_size=14,
sampling_ratio=2)
box_predictor = BBoxRegressor(2048,
num_classes=2,
bn_neck=cfg.MODEL.ROI_HEAD.BN_NECK)
roi_heads = SeqRoIHeads(
# OIM
num_pids=cfg.MODEL.LOSS.LUT_SIZE,
num_cq_size=cfg.MODEL.LOSS.CQ_SIZE,
oim_momentum=cfg.MODEL.LOSS.OIM_MOMENTUM,
oim_scalar=cfg.MODEL.LOSS.OIM_SCALAR,
# SeqNet
faster_rcnn_predictor=faster_rcnn_predictor,
reid_head=reid_head,
# parent class
box_roi_pool=box_roi_pool,
box_head=box_head,
box_predictor=box_predictor,
fg_iou_thresh=cfg.MODEL.ROI_HEAD.POS_THRESH_TRAIN,
bg_iou_thresh=cfg.MODEL.ROI_HEAD.NEG_THRESH_TRAIN,
batch_size_per_image=cfg.MODEL.ROI_HEAD.BATCH_SIZE_TRAIN,
positive_fraction=cfg.MODEL.ROI_HEAD.POS_FRAC_TRAIN,
bbox_reg_weights=None,
score_thresh=cfg.MODEL.ROI_HEAD.SCORE_THRESH_TEST,
nms_thresh=cfg.MODEL.ROI_HEAD.NMS_THRESH_TEST,
detections_per_img=cfg.MODEL.ROI_HEAD.DETECTIONS_PER_IMAGE_TEST,
)
transform = GeneralizedRCNNTransform(
min_size=cfg.INPUT.MIN_SIZE,
max_size=cfg.INPUT.MAX_SIZE,
image_mean=[0.485, 0.456, 0.406],
image_std=[0.229, 0.224, 0.225],
)
self.backbone = backbone
self.rpn = rpn
self.roi_heads = roi_heads
self.transform = transform
# loss weights
self.lw_rpn_reg = cfg.SOLVER.LW_RPN_REG
self.lw_rpn_cls = cfg.SOLVER.LW_RPN_CLS
self.lw_proposal_reg = cfg.SOLVER.LW_PROPOSAL_REG
self.lw_proposal_cls = cfg.SOLVER.LW_PROPOSAL_CLS
self.lw_box_reg = cfg.SOLVER.LW_BOX_REG
self.lw_box_cls = cfg.SOLVER.LW_BOX_CLS
self.lw_box_reid = cfg.SOLVER.LW_BOX_REID
import torch
import torchvision
from datasets.segmentation import CamVid
from models.deeplabv3plus import DeepLabV3plus
from models.resnet import build_resnet
encoder = build_resnet("18", output_stride=8)
model = DeepLabV3plus(encoder, 12)
dataset = CamVid("../seg_datasets/CamVid")
image1 = torchvision.transforms.functional.to_tensor(dataset[0][0])
image2 = torchvision.transforms.functional.to_tensor(dataset[1][0])
images = torch.stack([image1, image2], dim=0)
print("Without gradient checkpointing")
model(images, False)
print("\nWith gradient checkpointing")
model(images, True)