Menu
Total newbie here, I'm using this pytorch SegNet implementation with a '.pth' file containing weights from a 50 epochs training. How can I load a single test image and see the net prediction? I know this may sound like a stupid question but I'm stuck. What I've got is:
from segnet import SegNet
import torch
model = SegNet(2)
model.load_state_dict(torch.load('./model_segnet_epoch50.pth'))
How do I "use" the net on a single test picture?
345
A common PyTorch convention is to save these checkpoints using the . tar file extension. To load the models, first initialize the models and optimizers, then load the dictionary locally using torch. load() .
When loading a model on a GPU that was trained and saved on CPU, set the map_location argument in the torch. load() function to cuda:device_id . This loads the model to a given GPU device.
I provide with an example of ResNet152 pre-trained model.
def image_loader(loader, image_name):
image = Image.open(image_name)
image = loader(image).float()
image = torch.tensor(image, requires_grad=True)
image = image.unsqueeze(0)
return image
data_transforms = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor()
])
model_ft = models.resnet152(pretrained=True)
model_ft.eval()
print( np.argmax(model_ft(image_loader(data_transforms, $FILENAME)).detach().numpy()))
$FILENAME is the path and name of your image to be loaded. I got necessary help from this post.
94
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With
