在pytorch的torchvision中提供了许多深度学习领域各种经典的网络结构以及预训练模型,这些模型接口都被封装在torchvision.models中。
如图像分类的AlexNet、DenseNet、ResNet、MobileNet等
语义分割任务:DeepLabV3、FCN、LRASPP等以及目标检测任务:Faster R-CNN、FCOS、RetinaNet、SSD、SSDlite等。
对于这些模型都是已经训练好的,我们可以使用他们进行迁移学习,在此基础上根据需求对网络的结构进行修改,对于图像分类任务通常是修改最后的输出层,只需要把输出层的维数修改为我们目标任务即可。
resnet = torchvision.models.resnet152(pretrained=True) detection = torch.models.detection.maskrcnn_resnet50_fpn(pretrained=True) vit = torchvision.models.vision_transformer.vit_b_16(pretrained=True)