yearing1017 / global-wheat-detection Goto Github PK

# load a model; pre-trained on COCO（以下4句为pytorch官方教程例子）
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
num_classes = 2  # 1 class (wheat) + background
# 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)

for images, targets, image_ids in train_data_loader:
        
        images = list(image.to(device) for image in images)
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
        # 官方教程中写的是 model(images, targets):Returns losses and detections
        loss_dict = model(images, targets)  # Returns losses and detections