Coverage for src/deepdraw/configs/models/lwnet.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"""Little W-Net for image segmentation. 

6 

7The Little W-Net architecture contains roughly around 70k parameters and 

8closely matches (or outperforms) other more complex techniques. 

9 

10Reference: [GALDRAN-2020]_ 

11""" 

12 

13from torch.optim import Adam 

14from torch.optim.lr_scheduler import CosineAnnealingLR 

15 

16from deepdraw.models.losses import MultiWeightedBCELogitsLoss 

17from deepdraw.models.lwnet import lwnet 

18 

19# config 

20max_lr = 0.01 # start 

21min_lr = 1e-08 # valley 

22cycle = 50 # epochs for a complete scheduling cycle 

23 

24model = lwnet() 

25 

26criterion = MultiWeightedBCELogitsLoss() 

27 

28optimizer = Adam( 

29 model.parameters(), 

30 lr=max_lr, 

31) 

32 

33scheduler = CosineAnnealingLR( 

34 optimizer, 

35 T_max=cycle, 

36 eta_min=min_lr, 

37)