### Create / load model

# Faster - RCNN Model - pretrained on COCO
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
num_classes = 2

# 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)

Training the model doesn't use model.fit() function, it uses cycles.

    # let's train it for 10 epochs
    num_epochs = 10

    for epoch in range(num_epochs):
        # train for one epoch, printing every 10 iterations
        train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
        # update the learning rate
        # evaluate on the test dataset
        evaluate(model, data_loader_test, device=device)

How can I save the model?

  • 1
    $\begingroup$ That seems like Pytorch, not Keras. $\endgroup$
    – noe
    May 18, 2021 at 9:50
  • 2
    $\begingroup$ In Pytorch, it is done following the method here. $\endgroup$
    – Ubikuity
    May 18, 2021 at 14:20

1 Answer 1


PyTorch has a state_dict which stores the state of the model (in this case, the neural network) at any point in time. Saving it would involve dumping those states into a file which is easily done with:

torch.save(model.state_dict(), PATH)

When reloading the model, remember to first create the model class with its default weights and load the state dict from the file.

Here's a link to saving/loading pytorch modules (https://pytorch.org/tutorials/beginner/saving_loading_models.html).


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.