1#!/usr/bin/env python 

2# -*- coding: utf-8 -*- 

3 

4"""Little W-Net for image segmentation 

5 

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

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

8 

9Reference: [GALDRAN-2020]_ 

10""" 

11 

12from torch.optim import Adam 

13from torch.optim.lr_scheduler import CosineAnnealingLR 

14 

15from bob.ip.binseg.models.losses import MultiWeightedBCELogitsLoss 

16from bob.ip.binseg.models.lwnet import lwnet 

17 

18# config 

19max_lr = 0.01 # start 

20min_lr = 1e-08 # valley 

21cycle = 50 # epochs for a complete scheduling cycle 

22 

23model = lwnet() 

24 

25criterion = MultiWeightedBCELogitsLoss() 

26 

27optimizer = Adam( 

28 model.parameters(), 

29 lr=max_lr, 

30) 

31 

32scheduler = CosineAnnealingLR( 

33 optimizer, 

34 T_max=cycle, 

35 eta_min=min_lr, 

36)