pytorch的torchvision.models模块fasterrcnn目标检测之一

import torch
import torchvision.models as models
import torchvision

提前下载好权重文件,加载fasterrcnn模型

net = models.detection.fasterrcnn_resnet50_fpn(pretrained=False)
pthfile = r'G:\Anaconda3\envs\pytorch\Lib\site-packages\torchvision\models\fasterrcnn_resnet50_fpn_coco-258fb6c6.pth'
net.load_state_dict(torch.load(pthfile))
#print(net)

我们用PIL读取一张图片,进去看看

from PIL import Image
img = Image.open("maandgirl.jpg")
img

在这里插入图片描述

from torchvision import transforms
transform = transforms.Compose([            #[1]
 transforms.Resize(256),                    #[2]
 transforms.CenterCrop(224),                #[3]
 transforms.ToTensor(),                     #[4]
 transforms.Normalize(                      #[5]
 mean=[0.485, 0.456, 0.406],                #[6]
 std=
  • 0
    点赞
  • 12
    收藏
    觉得还不错? 一键收藏
  • 4
    评论
以下是class-agnostic模块作用于faster_RCNN目标检测PyTorch训练示例代码,注释已经加入到代码中: ``` python import torch import torchvision from torchvision.models.detection.faster_rcnn import FastRCNNPredictor from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor # 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 # 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 # (e.g., if your backbone produces a few feature maps of different sizes). anchor_generator = torchvision.models.detection.rpn.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 needs 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 = torchvision.models.detection.FasterRCNN(backbone, num_classes=2, rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler) # define a class-agnostic module class ClassAgnosticModule(torch.nn.Module): def __init__(self, in_channels, num_classes): super().__init__() self.conv = torch.nn.Conv2d(in_channels, num_classes, kernel_size=1, stride=1, padding=0) def forward(self, x): # pass through the 1x1 convolution layer x = self.conv(x) # flatten the tensor x = x.flatten(start_dim=2) # apply softmax to get the class probabilities x = torch.nn.functional.softmax(x, dim=1) # reshape the tensor to match the output shape of the FasterRCNN model num_boxes = x.shape[1] x = x.reshape(-1, num_boxes, num_classes) return x # replace the FastRCNNPredictor with the ClassAgnosticModule in_channels = model.roi_heads.box_predictor.cls_score.in_features num_classes = 2 model.roi_heads.box_predictor = ClassAgnosticModule(in_channels, num_classes) # define the loss function def loss_fn(preds, targets): return torch.nn.functional.cross_entropy(preds.squeeze(), targets) # define the optimizer optimizer = torch.optim.Adam(model.parameters(), lr=0.001) # define the data loader data_loader = torch.utils.data.DataLoader(dataset, batch_size=4, shuffle=True) # train the model for epoch in range(num_epochs): for images, targets in data_loader: # move the images and targets to the device images = list(image.to(device) for image in images) targets = [{k: v.to(device) for k, v in t.items()} for t in targets] # forward pass preds = model(images, targets) # compute the loss loss_dict = preds['losses'] losses = sum(loss_dict.values()) # backward pass optimizer.zero_grad() losses.backward() optimizer.step() ``` 该示例代码中,我们首先加载了一个预训练的分类模型,并删去了分类层。然后,我们定义了一个class-agnostic模块,并将FastRCNNPredictor替换为该模块。模型的其余部分与标准的FasterRCNN模型相同。最后,我们定义了一个损失函数和一个优化器,并使用数据加载器训练模型。 需要注意的是,该示例代码中的dataset和num_epochs变量没有给出,需要根据具体情况进行设置。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 4
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值