Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
4"""U-Net for image segmentation
6U-Net is a convolutional neural network that was developed for biomedical image
7segmentation at the Computer Science Department of the University of Freiburg,
8Germany. The network is based on the fully convolutional network (FCN) and its
9architecture was modified and extended to work with fewer training images and
10to yield more precise segmentations.
12Reference: [RONNEBERGER-2015]_
13"""
15from torch.optim.lr_scheduler import MultiStepLR
17from bob.ip.binseg.engine.adabound import AdaBound
18from bob.ip.binseg.models.losses import SoftJaccardBCELogitsLoss
19from bob.ip.binseg.models.unet import unet
21# config
22lr = 0.001
23betas = (0.9, 0.999)
24eps = 1e-08
25weight_decay = 0
26final_lr = 0.1
27gamma = 1e-3
28eps = 1e-8
29amsbound = False
31scheduler_milestones = [900]
32scheduler_gamma = 0.1
34model = unet()
36# optimizer
37optimizer = AdaBound(
38 model.parameters(),
39 lr=lr,
40 betas=betas,
41 final_lr=final_lr,
42 gamma=gamma,
43 eps=eps,
44 weight_decay=weight_decay,
45 amsbound=amsbound,
46)
48# criterion
49criterion = SoftJaccardBCELogitsLoss(alpha=0.7)
51# scheduler
52scheduler = MultiStepLR(
53 optimizer, milestones=scheduler_milestones, gamma=scheduler_gamma
54)