Coverage for src/deepdraw/configs/models/unet.py: 100%

1# SPDX-FileCopyrightText: Copyright © 2023 Idiap Research Institute <contact@idiap.ch> 

2# 

3# SPDX-License-Identifier: GPL-3.0-or-later 

4 

5"""U-Net for image segmentation. 

6 

7U-Net is a convolutional neural network that was developed for biomedical image 

8segmentation at the Computer Science Department of the University of Freiburg, 

9Germany. The network is based on the fully convolutional network (FCN) and its 

10architecture was modified and extended to work with fewer training images and 

11to yield more precise segmentations. 

12 

13Reference: [RONNEBERGER-2015]_ 

14""" 

15 

16from torch.optim.lr_scheduler import MultiStepLR 

17 

18from deepdraw.engine.adabound import AdaBound 

19from deepdraw.models.losses import SoftJaccardBCELogitsLoss 

20from deepdraw.models.unet import unet 

21 

22# config 

23lr = 0.001 

24betas = (0.9, 0.999) 

25eps = 1e-08 

26weight_decay = 0 

27final_lr = 0.1 

28gamma = 1e-3 

29eps = 1e-8 

30amsbound = False 

31 

32scheduler_milestones = [900] 

33scheduler_gamma = 0.1 

34 

35model = unet() 

36 

37# optimizer 

38optimizer = AdaBound( 

39 model.parameters(), 

40 lr=lr, 

41 betas=betas, 

42 final_lr=final_lr, 

43 gamma=gamma, 

44 eps=eps, 

45 weight_decay=weight_decay, 

46 amsbound=amsbound, 

47) 

48 

49# criterion 

50criterion = SoftJaccardBCELogitsLoss(alpha=0.7) 

51 

52# scheduler 

53scheduler = MultiStepLR( 

54 optimizer, milestones=scheduler_milestones, gamma=scheduler_gamma 

55)