FilesExpand file tree

main

/

optim.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

46 lines (36 loc) · 1.58 KB

main

/

optim.py

Copy path

Top

File metadata and controls

• Code
• Blame

46 lines (36 loc) · 1.58 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

import torch

import torch.nn

import utils.sam

import torch.optim

import torch.optim.lr_scheduler

import torch.optim.swa_utils

def get_optimizer_scheduler(model_name,

optim_name,

net,

lr,

momentum,

weight_decay,

max_epoch_cos = 200,

swa_lr = 0.05) :

## sgd + sam

sgd_optimizer = torch.optim.SGD(net.parameters(), lr=lr, momentum=momentum, weight_decay=weight_decay)

sam_sgd = utils.sam.SAM(net.parameters(), torch.optim.SGD, lr=lr, momentum=momentum, weight_decay=weight_decay)

## adamw + sam

adamw_optimizer = torch.optim.AdamW(net.parameters(), lr=lr, weight_decay=weight_decay)

sam_adamw = utils.sam.SAM(net.parameters(), torch.optim.AdamW, lr=lr, weight_decay=weight_decay)

## convmixer uses adamw optimzer while cnn backbones uses sgd

if model_name in ["convmixer", "vit_cifar"] :

if optim_name in ['sam', 'fmfp'] :

optimizer = sam_adamw

else :

optimizer = adamw_optimizer

else:

if optim_name in ['sam', 'fmfp'] :

optimizer = sam_sgd

else :

optimizer = sgd_optimizer

cos_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_epoch_cos)

## swa model

swa_model = torch.optim.swa_utils.AveragedModel(net)

swa_scheduler = torch.optim.swa_utils.SWALR(optimizer, swa_lr=swa_lr)

return optimizer, cos_scheduler, swa_model, swa_scheduler