saving checkpoints in "model.pth" format and visualizing the segmentation results

LeeJunHyun / Image_Segmentation

Pytorch implementation of U-Net, R2U-Net, Attention U-Net, and Attention R2U-Net.

2.66k stars 594 forks source link

saving checkpoints in "model.pth" format and visualizing the segmentation results #78

Closed k-nayak closed 1 year ago

k-nayak commented 2 years ago

Hello @LeeJunHyun , thank you for the amazing code. I want to use Attention Unet for real-time segmentation and see how well it works. I need weight files for using the model with a camera setup usually in ".pth" format so i can feed new images and see the segmentation results. Is there any way to save the checkpoint or is the ".pkl" something that can be used instead. Thanks in advance.

Lloyd-Pottiger commented 2 years ago

The Pytorch offical tutorials may can help you https://pytorch.org/tutorials/beginner/basics/saveloadrun_tutorial.html, and you can just change the unet_path with '.pth'.