def __init__(self, num_classes):
super().__init__()
model = torchvision.models.detection.maskrcnn_resnet50_fpn(
pretrained=True)
self.transform = model.transform
self.backbone = model.backbone
self.classifier = DeepLabHead(model.backbone.out_channels, num_classes)
def deeplabv3_efficientnet(n):
backbone = EfficientNetFeatureExtractor(n)
inplanes = efficientnet_inplanes[n]
classifier = DeepLabHead(inplanes, 3)
return SimpleSegmentationModel(backbone, classifier)
def __init__(self, num_classes, backbone_fn, chip_size=224):
super().__init__()
if getattr(backbone_fn, '_is_multispectral', False):
self.backbone = create_body(backbone_fn,
pretrained=True,
cut=_get_backbone_meta(
backbone_fn.__name__)['cut'])
else:
self.backbone = create_body(backbone_fn, pretrained=True)
backbone_name = backbone_fn.__name__
## Support for different backbones
if "densenet" in backbone_name or "vgg" in backbone_name:
hookable_modules = list(self.backbone.children())[0]
else:
hookable_modules = list(self.backbone.children())
if "vgg" in backbone_name:
modify_dilation_index = -5
else:
modify_dilation_index = -2
if backbone_name == 'resnet18' or backbone_name == 'resnet34':
module_to_check = 'conv'
else:
module_to_check = 'conv2'
## Hook at the index where we need to get the auxillary logits out
self.hook = hook_output(hookable_modules[modify_dilation_index])
custom_idx = 0
for i, module in enumerate(hookable_modules[modify_dilation_index:]):
dilation = 2 * (i + 1)
padding = 2 * (i + 1)
for n, m in module.named_modules():
if module_to_check in n:
m.dilation, m.padding, m.stride = (dilation, dilation), (
padding, padding), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
if "vgg" in backbone_fn.__name__:
if isinstance(module, nn.Conv2d):
dilation = 2 * (custom_idx + 1)
padding = 2 * (custom_idx + 1)
module.dilation, module.padding, module.stride = (
dilation, dilation), (padding, padding), (1, 1)
custom_idx += 1
## returns the size of various activations
feature_sizes = model_sizes(self.backbone, size=(chip_size, chip_size))
## Geting the number of channel persent in stored activation inside of the hook
num_channels_aux_classifier = self.hook.stored.shape[1]
## Get number of channels in the last layer
num_channels_classifier = feature_sizes[-1][1]
self.classifier = DeepLabHead(num_channels_classifier, num_classes)
self.aux_classifier = FCNHead(num_channels_aux_classifier, num_classes)
def _get_model():
# DeepLab V3 encoder, ResNet 101 decoder, use pretrained model.
model = models.segmentation.deeplabv3_resnet101(pretrained=True,
progress=True)
# Feature vector size = 2048 (resnet101), output channels = 1
model.classifier = DeepLabHead(2048, 1)
return model
def __init__(self, num_classes):
super(DeepSqueeze, self).__init__()
backbone = squeeze.features
classifier = DeepLabHead(512, num_classes)
self.backbone = backbone
self.classifier = classifier
self.aux_classifier = None
def _create_deeplab(num_class, backbone, pretrained=True, **kwargs):
'''
Create default torchvision pretrained model with resnet101.
'''
model = models.segmentation.deeplabv3_resnet101(pretrained=True, progress=True, **kwargs)
model = _DeepLabOverride(model.backbone, model.classifier, model.aux_classifier)
model.classifier = DeepLabHead(2048, num_class)
model.aux_classifier = FCNHead(1024, num_class)
return model
def efficient_deeplabv3(args):
# args.backbone には 'efficientnet-b0'などが入る
efficientnet = EfficientNetExtractor.from_name(args.backbone,
in_channels=1)
num_ch = efficientnet._conv_head.out_channels
classifier = DeepLabHead(num_ch, args.num_classes)
base_model = Deeplabv3
model = base_model(efficientnet, classifier)
return model
def createDeepLabv3(outputchannels=1):
"""DeepLabv3 class with custom head
Args:
outputchannels (int, optional): The number of output channels
in your dataset masks. Defaults to 1.
Returns:
model: Returns the DeepLabv3 model with the ResNet50 backbone.
"""
model = models.segmentation.deeplabv3_resnet50(pretrained=True,
progress=True)
model.classifier = DeepLabHead(2048, outputchannels)
return model
def __init__(self, in_chan=3, out_chan=2, pretrained=False):
super(densenet50, self).__init__()
self.model = torchvision.models.segmentation.deeplabv3_resnet50(
pretrained=False, pretrained_backbone=pretrained)
self.model.classifier = DeepLabHead(2048, out_chan)
if in_chan != 3:
self.model.backbone.conv1 = nn.Conv2d(in_chan,
64,
kernel_size=7,
stride=2,
padding=3,
bias=False)
def deeplabv3_resnet101():
backbone = resnet.__dict__['resnet101'](
pretrained=True, replace_stride_with_dilation=[False, True, True])
return_layers = {'layer4': 'out'}
backbone = IntermediateLayerGetter(backbone, return_layers=return_layers)
inplanes = 2048
classifier = DeepLabHead(inplanes, 3)
return SimpleSegmentationModel(backbone, classifier)
def createDeepLabv3(outputchannels, backboneFreez=False):
model = models.segmentation.deeplabv3_resnet101(pretrained=True, progress=True)
if backboneFreez:
for param in model.parameters():
param.requires_grad = False
# Allocate a new classifier
model.classifier = DeepLabHead(2048, outputchannels)
# Set the model in training mode
model.train()
return model
def get_model(num_classes):
# load an instance segmentation model pre-trained pre-trained on COCO
model = deeplabv3_resnet101(pretrained=True)
for param in model.parameters():
param.requires_grad = False
deeplab_head_in_channels = model.classifier[0].convs[0][0].in_channels
fcn_head_in_channels = model.classifier[1].in_channels
model.classifier = DeepLabHead(deeplab_head_in_channels, num_classes)
return model
def __init__(self, num_classes):
super().__init__()
self.num_classes = num_classes
self.segm = deeplabv3_resnet101(pretrained=True)
for param in self.segm.parameters():
param.requires_grad = False
self.segm.classifier = DeepLabHead(in_channels=2048,
num_classes=num_classes)
self.segm.classifier.apply(_init_weights)
def __loadSegModel(self):
seg_model = models.segmentation.deeplabv3_resnet101(pretrained=True,
progress=True)
seg_model.classifier = DeepLabHead(2048, 1)
checkpoint = torch.load(self.seg_checkpoint_path,
map_location=self.device)
seg_model.load_state_dict(checkpoint['model_state_dict'])
seg_model.eval()
return seg_model
def createDeepLabv3(outputchannels=1, mode='train'):
"""DeepLabv3 class with custom head
Args:
outputchannels (int, optional): The number of output channels
in your dataset masks. Defaults to 1.
Returns:
model: Returns the DeepLabv3 model with the ResNet101 backbone.
"""
model = models.segmentation.deeplabv3_resnet50(pretrained=True,
progress=True)
model.classifier = DeepLabHead(2048, outputchannels)
# Set the model in training mode
if mode == 'train':
model.train()
else:
model.eval()
return model.float()
def deeplabv3(pretrained=False,
resnet="res103",
head_in_ch=2048,
num_classes=21):
resnet = {"res53": resnet53, "res103": resnet103}[resnet]
net = SmallDeepLab(backbone=IntermediateLayerGetter(resnet(
pretrained=True, replace_stride_with_dilation=[False, True, True]),
return_layers={
'layer2': 'res2',
'layer4': 'out'
}),
classifier=DeepLabHead(head_in_ch, num_classes))
if pretrained:
state_dict = load_state_dict_from_url(
'https://download.pytorch.org/models/deeplabv3_resnet101_coco-586e9e4e.pth',
progress=True)
net.load_state_dict(state_dict)
return net
def pre_deeplabv3_resnet101(in_channel, out_channel):
model = deeplabv3_resnet101(pretrained=False)
url = 'https://download.pytorch.org/models/deeplabv3_resnet101_coco-586e9e4e.pth'
model_dict = model.state_dict()
pretrained_dict = model_zoo.load_url(url, progress=False)
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
model.backbone.conv1 = nn.Conv2d(in_channel,
64,
kernel_size=7,
stride=2,
padding=3,
bias=False)
model.classifier = DeepLabHead(2048, out_channel)
return model
def deeplabv3(pretrained=False,
resnet="res103",
head_in_ch=2048,
num_classes=21):
resnet = {
"res53": resnet53,
"res103": resnet103,
"res50": resnet50,
"res101": resnet101
}[resnet]
net = SmallDeepLab(backbone=IntermediateLayerGetter(resnet(
pretrained=True, replace_stride_with_dilation=[False, True, True]),
return_layers={
'layer2': 'res2',
'layer4': 'out'
}),
classifier=DeepLabHead(head_in_ch, num_classes))
return net
def createDeepLabv3(outputchannels=1):
"""DeepLabv3 class with custom head
Args:
outputchannels (int, optional): The number of output channels
in your dataset masks. Defaults to 1.
Returns:
model: Returns the DeepLabv3 model with the ResNet101 backbone.
"""
model = models.segmentation.deeplabv3_resnet101(pretrained=False,
progress=True)
model.backbone.conv1 = nn.Conv2d(2,
64,
kernel_size=(7, 7),
stride=(2, 2),
padding=(3, 3),
bias=False)
model.classifier = DeepLabHead(2048, outputchannels)
# Set the model in training mode
# model.train()
return model
def predict(model_checkpoint, image_path, out_file):
"""
Args:
model_checkpoint (string): path to model checkpoint
image_path (string): path to an aerial image
out_file (string): path to save the prediction mask.
"""
DEVICE = torch.device("cpu")
RGB_MEAN = [0.485, 0.456, 0.406]
RGB_STD = [0.229, 0.224, 0.225]
image = Image.open(str(image_path))
image = image.resize((512, 512), Image.BILINEAR)
image_transforms = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize(RGB_MEAN, RGB_STD)])
image_tensor = image_transforms(image)[None]
# model = FastSCNN(num_classes=3)
# model = UNet(n_channels=3, n_classes=3, bilinear=True)
model = deeplabv3(pretrained=False, progress=True)
model.classifier = DeepLabHead(2048, 3)
model.eval()
model.to(DEVICE)
image_tensor = image_tensor.to(device=DEVICE, dtype=torch.float)
model.load_state_dict(
torch.load(model_checkpoint,
map_location=lambda storage, loc: storage))
out = model(image_tensor)['out'][0].squeeze()
out_max = out.max(0, keepdim=False)[1].cpu().numpy()
final_image = np.zeros((out_max.shape[0], out_max.shape[1], 3),
dtype=np.int)
final_image[(out_max == 0), :] = np.array([255, 255, 255])
final_image[(out_max == 1), :] = np.array([255, 0, 0])
final_image[(out_max == 2), :] = np.array([0, 0, 255])
# image.show()
final_image_pil = Image.fromarray(np.uint8(final_image))
final_image_pil.show()
final_image_pil.save(out_file)
def __init__(self, model_weight, img_dir, out_dir, device):
self.model = models.segmentation.deeplabv3_resnet50(pretrained=False,
progress=True)
self.model.classifier = DeepLabHead(2048, 1)
self.model.to(device)
model_weight = torch.load(model_weight)
model_weight = model_weight["model_state_dict"]
self.model.load_state_dict(model_weight, strict=False)
self.img = Prediction(img_dir=img_dir)
self.dataloader = DataLoader(self.img,
batch_size=16,
num_workers=4,
shuffle=False)
self.out_dir = out_dir
os.makedirs(self.out_dir, exist_ok=True)
self.device = device
def main():
nlabel = 6
df = pd.read_csv(os.path.join(args.root, "train.csv"))
dataset = ProstateSeg(df, args.root, 2048)
loader = DataLoader(dataset,num_workers = 4, pin_memory=True)
model = models.segmentation.deeplabv3_resnet101(
pretrained=True, progress=True)
model.classifier = DeepLabHead(2048, nlabel)
print('Evaluate using {}...'
.format(args.checkpoint))
weight_file=os.path.join(args.checkpoint)
model.load_state_dict(torch.load(weight_file, map_location=lambda storage, loc: storage))
if not os.path.exists(args.dump):
os.mkdir(args.dump)
model.eval()
model.to(device)
with torch.no_grad():
t0 = time.time()
for i, inputs in enumerate(loader):
t1 = time.time()
imid = df.iloc[i].image_id
provider = df.iloc[i].data_provider
grade = df.iloc[i].isup_grade
score = df.iloc[i].gleason_score
inputs = inputs.to(device)
output = model(inputs)
pred = output['out'].argmax(dim=1).detach().cpu()
pp = np.zeros(nlabel)
npix = np.prod(pred.shape)
for i in range(nlabel):
pp[i] = pred.eq(i).sum().item() / npix
pp = pp[1:]/(pp[1:].sum())
t2 = time.time()
print("{:.3f}\t{:.3f}\t{:.3f}\t{:.3f}\t{:.3f}|{},{},{},{},{:.1f}s/{:.1f}s".format(
*pp, grade, score, provider, imid, t2-t1, (t2-t0)/(i+1)))
torch.save(pred, os.path.join(args.dump, imid))
def fit_with(lr: float, wd: float):
model = torchvision.models.segmentation.deeplabv3_resnet50(
pretrained=True)
model.classifier = DeepLabHead(2048, 2)
model.to(device)
criterion = torch.nn.CrossEntropyLoss().to(device)
optimizer = torch.optim.Adam(model.parameters(),
lr=lr,
weight_decay=wd)
model, acc, loss = train_model(model=model,
criterion=criterion,
dls=dataloaders,
opt=optimizer,
n_classes=2,
name=f'h_param_{lr:.5f}_{wd:.5f}',
log_path=args.logs_dir,
epochs=args.epochs,
verbose=False)
return acc
def __init__(self, num_classes: int = 2, ignore_index: Optional[int] = None, lr: float = 0.001,
weight_decay: float = 0.001, aux_loss_factor: float = 0.3):
super().__init__()
self.num_classes = num_classes
self.ignore_index = ignore_index
self.lr = lr
self.weight_decay = weight_decay
self.aux_loss_factor = aux_loss_factor
# Create model from pre-trained DeepLabv3
self.model = deeplabv3_resnet101(pretrained=True, progress=True)
self.model.aux_classifier = FCNHead(1024, self.num_classes)
self.model.classifier = DeepLabHead(2048, self.num_classes)
# Setup trainable layers
self.model.requires_grad_(True)
self.model.backbone.requires_grad_(False)
# Loss function and metrics
self.focal_tversky_loss = FocalTverskyMetric(
self.num_classes,
alpha=0.7,
beta=0.3,
gamma=4.0 / 3.0,
ignore_index=self.ignore_index,
)
self.accuracy_metric = Accuracy(ignore_index=self.ignore_index)
self.iou_metric = JaccardIndex(
num_classes=self.num_classes,
reduction="none",
ignore_index=self.ignore_index,
)
self.precision_metric = Precision(num_classes=self.num_classes, ignore_index=self.ignore_index,
average='weighted', mdmc_average='global')
self.recall_metric = Recall(num_classes=self.num_classes, ignore_index=self.ignore_index,
average='weighted', mdmc_average='global')
def main():
# dataset
train_df = pd.read_csv(os.path.join(args.root, "train.csv"))
is_radboud = (train_df['data_provider'] == 'radboud')
by_radboud = train_df[is_radboud]
#by_karolinska = train_df[np.logical_not(is_radboud)]
num = 0
for idx, row in by_radboud.iterrows():
img_id = row['image_id']
mask_path = os.path.join(args.root, "train_label_masks",
img_id + "_mask.tiff")
if not os.path.isfile(mask_path):
num += 1
#print("{}:{} mask not work!".format(num,img_id))
by_radboud = by_radboud.drop(idx)
df = {}
df['train'], df['val'] = train_test_split(by_radboud,
stratify=by_radboud.isup_grade,
test_size=20,
random_state=42)
dataset = {
'val':
ProstateSeg(df['val'], args.root, args.size,
(args.crop_size, args.crop_size), 'val')
}
loader = {
'val':
DataLoader(dataset['val'], num_workers=args.workers, pin_memory=True)
}
model = models.segmentation.deeplabv3_resnet101(pretrained=True,
progress=True)
model.classifier = DeepLabHead(2048, nlabel)
if args.checkpoint:
print('Resuming training from epoch {}, loading {}...'.format(
args.resume_epoch, args.checkpoint))
weight_file = os.path.join(args.output_folder, args.checkpoint)
model.load_state_dict(
torch.load(weight_file, map_location=lambda storage, loc: storage))
model.to(device)
criterion = FocalLoss(alpha=torch.tensor(
[1, 1.4, 6, 6, 5, 8], dtype=torch.float32, device=device))
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=args.weight_decay)
for epoch in range(args.resume_epoch, args.epochs):
s1, s2 = int(0.9 * args.crop_size), int(1.1 * args.crop_size)
crop = (randint(s1, s2), randint(s1, s2))
dataset['train'] = ProstateSeg(df['train'], args.root, args.size, crop,
'train')
loader['train'] = DataLoader(dataset['train'],
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
adjust_lr(optimizer, epoch, args)
for phase in ['train', 'val']:
t0 = time.time()
print("========={}:{}=========".format(phase, epoch))
if phase == 'train':
model.train()
else:
model.eval()
num = 0
correct = 0
nums = np.zeros(6, dtype=int)
pros = np.zeros(6, dtype=int)
corrects = np.zeros(6, dtype=int)
for i, (inputs, masks) in enumerate(loader[phase]):
t1 = time.time()
if i == 0: print(inputs.shape)
inputs = inputs.to(device)
masks = masks.to(device)
optimizer.zero_grad()
with torch.set_grad_enabled(phase == 'train'):
output = model(inputs)
loss = criterion(output['out'], masks)
if phase == 'train':
loss.backward()
optimizer.step()
pred = output['out'].argmax(dim=1)
correct = pred.eq(masks).sum().item()
npixel = np.prod(masks.shape)
acc = correct * 100 / npixel
num += masks.size(0)
if (i + 1) % args.log == 0:
t2 = time.time()
s = "({},{:.1f}s,{:.1f}s) Loss:{:.3f} Acc:{:.3f}"
print(
s.format(num, t2 - t1, (t2 - t0) / (i + 1),
loss.item(), acc))
if phase == 'val':
for i in range(nlabel):
t = masks.eq(i)
p = pred.eq(i)
nums[i] += t.sum().item()
pros[i] += p.sum().item()
corrects[i] += (p & t).sum().item()
if epoch % 1 == 0 and phase == "train":
torch.save(
model.state_dict(),
os.path.join(args.output_folder,
"deeplab-{}.pth".format(epoch)))
if phase == 'val':
print("recall:" + "|".join(
["{:.5f}".format(c / n) for c, n in zip(corrects, nums)]))
print("precision:" + "|".join(
["{:.5f}".format(c / p) for c, p in zip(corrects, pros)]))
def __init__(self):
super().__init__()
self.model = deeplabv3(pretrained=False, progress=True)
self.model.classifier = DeepLabHead(2048, 3)
def loadModel(model_arch="",
classes=None,
pre_trained_path=None,
expType=None,
trainable_backbone_flag=False,
lower_features=False):
print("Load model architecture ... ")
if (model_arch == "deeplabv3_resnet101_orig"):
print("deeplab_resnet architecture selected ...")
model = models.segmentation.deeplabv3_resnet101(pretrained=True,
progress=True)
for params in model.parameters():
params.requires_grad = trainable_backbone_flag
model.classifier[-1] = torch.nn.Conv2d(256,
len(classes),
kernel_size=(1, 1))
model.aux_classifier[-1] = torch.nn.Conv2d(256,
len(classes),
kernel_size=(1, 1))
features = model.backbone
if (pre_trained_path != None):
print("load pre-trained-weights ... ")
model_dict_state = torch.load(
pre_trained_path) # + "/best_model.pth")
model.load_state_dict(model_dict_state['net'])
return model, features
elif (model_arch == "fcn_resnet101_orig"):
print("deeplab_resnet architecture selected ...")
model = models.segmentation.fcn_resnet101(pretrained=True,
progress=True)
for params in model.parameters():
params.requires_grad = trainable_backbone_flag
model.classifier[-1] = torch.nn.Conv2d(512,
len(classes),
kernel_size=(1, 1))
model.aux_classifier[-1] = torch.nn.Conv2d(256,
len(classes),
kernel_size=(1, 1))
features = model.backbone
if (pre_trained_path != None):
print("load pre-trained-weights ... ")
model_dict_state = torch.load(
pre_trained_path) # + "/best_model.pth")
model.load_state_dict(model_dict_state['net'])
return model, features
elif (model_arch == "deeplabv3_resnet101"):
print("deeplabv3_resnet101 architecture selected ...")
backbone_net = CNN(model_arch="resnet101",
n_classes=len(classes),
include_top=False,
pretrained=trainable_backbone_flag,
lower_features=lower_features)
if (lower_features == True):
classifier = nn.Sequential(DeepLabHead(256, len(classes)),
# nn.Softmax()
)
else:
classifier = nn.Sequential(DeepLabHead(2048, len(classes)),
# nn.Softmax()
)
features = backbone_net
model = models.segmentation.DeepLabV3(backbone=backbone_net,
classifier=classifier,
aux_classifier=None)
if (pre_trained_path != None):
print("load pre-trained-weights ... ")
model_dict_state = torch.load(
pre_trained_path) # + "/best_model.pth")
model.load_state_dict(model_dict_state['net'])
return model, features
elif (model_arch == "deeplabv3_vgg16"):
print("deeplabv3_vgg architecture selected ...")
# backbone_net = CNN(model_arch="resnet101", n_classes=len(classes), include_top=False)
backbone_net = CNN(model_arch="vgg16",
n_classes=len(classes),
include_top=False,
pretrained=trainable_backbone_flag,
lower_features=lower_features)
if (lower_features == True):
classifier = nn.Sequential(DeepLabHead(64, len(classes)),
# nn.Softmax()
)
else:
classifier = nn.Sequential(DeepLabHead(512, len(classes)),
# nn.Softmax()
)
features = backbone_net
model = models.segmentation.DeepLabV3(backbone=backbone_net,
classifier=classifier,
aux_classifier=None)
#print(model)
#exit()
if (pre_trained_path != None):
print("load pre-trained-weights ... ")
model_dict_state = torch.load(
pre_trained_path) # + "/best_model.pth")
model.load_state_dict(model_dict_state['net'])
# Find total parameters and trainable parameters
total_params = sum(p.numel() for p in model.parameters())
print("total_params:" + str(total_params))
total_trainable_params = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print("total_trainable_params: " + str(total_trainable_params))
#exit()
return model, features
elif (model_arch == "deeplabv3_mobilenet"):
print("deeplabv3_mobilenet architecture selected ...")
backbone_net = CNN(model_arch="mobilenet",
n_classes=len(classes),
include_top=False,
pretrained=trainable_backbone_flag,
lower_features=lower_features)
if (lower_features == True):
classifier = nn.Sequential(DeepLabHead(32, len(classes)),
# nn.Softmax()
)
else:
classifier = nn.Sequential(DeepLabHead(1280, len(classes)),
# nn.Softmax()
)
features = backbone_net
model = models.segmentation.DeepLabV3(backbone=backbone_net,
classifier=classifier,
aux_classifier=None)
if (pre_trained_path != None):
print("load pre-trained-weights ... ")
model_dict_state = torch.load(pre_trained_path)
model.load_state_dict(model_dict_state['net'])
return model, features
elif (model_arch == "deeplabv3_squeezenet"):
print("deeplabv3_mobilenet architecture selected ...")
backbone_net = CNN(model_arch="squeezenet",
n_classes=len(classes),
include_top=False,
pretrained=trainable_backbone_flag,
lower_features=lower_features)
if (lower_features == True):
classifier = nn.Sequential(DeepLabHead(128, len(classes)),
# nn.Softmax()
)
else:
classifier = nn.Sequential(DeepLabHead(512, len(classes)),
# nn.Softmax()
)
features = backbone_net
model = models.segmentation.DeepLabV3(backbone=backbone_net,
classifier=classifier,
aux_classifier=None)
if (pre_trained_path != None):
print("load pre-trained-weights ... ")
model_dict_state = torch.load(
pre_trained_path) # + "/best_model.pth")
model.load_state_dict(model_dict_state['net'])
return model, features
elif (model_arch == "fcn_vgg16"):
print("fcn_vgg16 architecture selected ...")
backbone_net = CNN(model_arch="vgg16",
n_classes=len(classes),
include_top=False,
pretrained=trainable_backbone_flag,
lower_features=lower_features)
if (lower_features == True):
classifier = nn.Sequential(FCNHead(64, len(classes)),
# nn.Softmax()
)
else:
classifier = nn.Sequential(FCNHead(512, len(classes)),
# nn.Softmax()
)
features = backbone_net
model = models.segmentation.FCN(backbone=backbone_net,
classifier=classifier,
aux_classifier=None)
# print(model)
if (pre_trained_path != None):
print("load pre-trained-weights ... ")
model_dict_state = torch.load(
pre_trained_path) # + "/best_model.pth")
model.load_state_dict(model_dict_state['net'])
return model, features
elif (model_arch == "fcn_resnet101"):
print("fcn_resnet101 architecture selected ...")
backbone_net = CNN(model_arch="resnet101",
n_classes=len(classes),
include_top=False,
pretrained=trainable_backbone_flag,
lower_features=lower_features)
if (lower_features == True):
classifier = nn.Sequential(FCNHead(256, len(classes)),
# nn.Softmax()
)
else:
classifier = nn.Sequential(FCNHead(2048, len(classes)),
# nn.Softmax()
)
features = backbone_net
model = models.segmentation.FCN(backbone=backbone_net,
classifier=classifier,
aux_classifier=None)
if (pre_trained_path != None):
print("load pre-trained-weights ... ")
model_dict_state = torch.load(
pre_trained_path) # + "/best_model.pth")
model.load_state_dict(model_dict_state['net'])
# Find total parameters and trainable parameters
total_params = sum(p.numel() for p in model.parameters())
print("total_params:" + str(total_params))
total_trainable_params = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print("total_trainable_params: " + str(total_trainable_params))
#exit()
return model, features
elif (model_arch == "fcn_squeezenet"):
print("deeplabv3_squeezenet architecture selected ...")
backbone_net = CNN(model_arch="squeezenet",
n_classes=len(classes),
include_top=False,
pretrained=trainable_backbone_flag,
lower_features=lower_features)
if (lower_features == True):
classifier = nn.Sequential(FCNHead(128, len(classes)),
# nn.Softmax()
)
else:
classifier = nn.Sequential(FCNHead(512, len(classes)),
# nn.Softmax()
)
features = backbone_net
model = models.segmentation.FCN(backbone=backbone_net,
classifier=classifier,
aux_classifier=None)
if (pre_trained_path != None):
print("load pre-trained-weights ... ")
model_dict_state = torch.load(
pre_trained_path) # + "/best_model.pth")
model.load_state_dict(model_dict_state['net'])
# Find total parameters and trainable parameters
total_params = sum(p.numel() for p in model.parameters())
print("total_params:" + str(total_params))
total_trainable_params = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print("total_trainable_params: " + str(total_trainable_params))
# exit()
return model, features
elif (model_arch == "fcn_mobilenet"):
print("deeplabv3_mobilenet architecture selected ...")
backbone_net = CNN(model_arch="mobilenet",
n_classes=len(classes),
include_top=False,
pretrained=trainable_backbone_flag,
lower_features=lower_features)
if (lower_features == True):
classifier = nn.Sequential(FCNHead(32, len(classes)),
# nn.Softmax()
)
else:
classifier = nn.Sequential(FCNHead(1280, len(classes)),
# nn.Softmax()
)
features = backbone_net
model = models.segmentation.FCN(backbone=backbone_net,
classifier=classifier,
aux_classifier=None)
if (pre_trained_path != None):
print("load pre-trained-weights ... ")
model_dict_state = torch.load(
pre_trained_path) # + "/best_model.pth")
model.load_state_dict(model_dict_state['net'])
# Find total parameters and trainable parameters
total_params = sum(p.numel() for p in model.parameters())
print("total_params:" + str(total_params))
total_trainable_params = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print("total_trainable_params: " + str(total_trainable_params))
# exit()
return model, features
else:
print("ERROR: select valid model architecture!")
exit()
def __init__(self, num_classes=12):
super(TorchVisionDeepLabv3_ResNet50, self).__init__()
self.seg_model = torchvision.models.segmentation.deeplabv3_resnet50(pretrained=True)
self.seg_model.classifier = DeepLabHead(2048, num_classes)
kwargs = {}
if args.num_classes and not segmentation:
logging.info("Using num_classes = %d" % args.num_classes)
kwargs["num_classes"] = args.num_classes
model = model_classes[mname](pretrained=args.pretrained, progress=args.verbose, **kwargs)
if args.extract != None:
logging.info("Extract layers: " + ", ".join(args.extract))
return_nodes = {
layer: layer for layer in args.extract
}
model = create_feature_extractor(model, return_nodes=return_nodes)
if segmentation and 'deeplabv3' in mname:
model.classifier = DeepLabHead(2048, args.num_classes)
if args.to_dd_native:
# Make model NativeModuleWrapper compliant
model = Wrapper(model)
model.eval()
# tracing or scripting model (default)
if args.trace:
example = get_image_input(args.batch_size, args.img_width, args.img_height)
script_module = torch.jit.trace(model, example)
else:
script_module = torch.jit.script(model)
filename = os.path.join(
"""
# Set the number of output channels (= number of classes)
num_classes = 7
# Model definition
seg_model = models.segmentation.deeplabv3_resnet101(pretrained=True,
progress=True)
# The auxiliary classifier is removed, and the pretrained weights are frozen
seg_model.aux_classifier = None
for param in seg_model.parameters():
param.requires_grad = False
# The pretrained classifier is replaced by a new one with a custom number of classes.
# Since it comes after the freeze, its weights won't be frozen. They are the ones that we will fine-tune.
seg_model.classifier = DeepLabHead(2048, num_classes)
# Model serialisation
model_filename = 'custom_segmented.pt'
pathToModel = pathlib.Path.cwd().joinpath('..', 'models', model_filename)
print('File name for saved model: ', pathToModel)
# Loss function
criterion = torch.nn.CrossEntropyLoss(
) # combines nn.LogSoftmax() and nn.NLLLoss(), well suited for multiclass classification problems
# Optimizer definition
optimizer = torch.optim.SGD(seg_model.parameters(), lr=0.001, momentum=0.9)
# Use cpu/gpu based on availability
seg_model.to(device)