Skip to content

Instantly share code, notes, and snippets.

@magesh-technovator

magesh-technovator/load_model.py

Last active May 20, 2021 10:40
Show Gist options
  • Save magesh-technovator/934df1ddd9a32b041d1815bd1faca17e to your computer and use it in GitHub Desktop.
Save magesh-technovator/934df1ddd9a32b041d1815bd1faca17e to your computer and use it in GitHub Desktop.
Download ZIP
Load DeepLabV3 model
Raw
load_model.py
from torchvision import models
from torchvision.models.segmentation.deeplabv3 import DeepLabHead
def createDeepLabv3(outputchannels=1):
model = models.segmentation.deeplabv3_resnet101(
pretrained=True, progress=True)
# Added a Tanh activation after the last convolution layer
model.classifier = DeepLabHead(2048, outputchannels)
# Set the model in training mode
model.train()
return model
model = createDeepLabv3(3)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment