How to load a trained ML model with pytorch

Yiwei Miao 0
Viewed 21

I am very new to Machine Learning, and very lack of the experience. Recently i am trying to collect my own datasets and see the performance of it on a trained model. In README.md it only tells me how to train my own model, but it did not tell me how to load his model direct from here.

For instance, i want to try this trained model (culane_r18.pth.zip). How i am suppose to load this model? Should i start a new python file?

# What should i add here?
... 

FILE = "culane_r18.pth"

device = torch.device('cpu')
model = Model(*args, **kwargs)
model.load_state_dict(torch.load(FILE, map_location=device))

...

I have only finished the pytorch tutorial, and now i do not know what should i do with is current model.

Can anyone give me some help or tell me where should i find more details about this topic.

pytorch
1 Answers

I didn't try testing it but personally I'd add another function based on Runner.validate() here so that you don't need to bother with cfg. Maybe like

    def predict(self, data):
        self.net.eval()
        with torch.no_grad():
            output = self.net(data)
            output = self.net.module.heads.get_lanes(output)
        return output

Which you can just call runner.predict(data) similar to how validate() is called in main.py here.

You should be able to find the appropriate dataset from here and dataloader from here.

Natthaphon Hongcharoen0
Related