File size: 3,394 Bytes
d7e58f0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
import json
import torch
import torch.nn as nn


def match_name_keywords(n: str, name_keywords: list):
    out = False
    for b in name_keywords:
        if b in n:
            out = True
            break
    return out


def get_param_dict(args, model_without_ddp: nn.Module):
    try:
        param_dict_type = args.param_dict_type
    except:
        param_dict_type = 'default'
    assert param_dict_type in ['default', 'ddetr_in_mmdet', 'large_wd']

    # by default
    if param_dict_type == 'default':
        param_dicts = [{
            'params': [
                p for n, p in model_without_ddp.named_parameters()
                if 'backbone' not in n and p.requires_grad
            ]
        }, {
            'params': [
                p for n, p in model_without_ddp.named_parameters()
                if 'backbone' in n and p.requires_grad
            ],
            'lr':
            args.lr_backbone,
        }]
        return param_dicts

    if param_dict_type == 'ddetr_in_mmdet':
        param_dicts = [{
            'params': [
                p for n, p in model_without_ddp.named_parameters()
                if not match_name_keywords(n, args.lr_backbone_names)
                and not match_name_keywords(n, args.lr_linear_proj_names)
                and p.requires_grad
            ],
            'lr':
            args.lr,
        }, {
            'params': [
                p for n, p in model_without_ddp.named_parameters()
                if match_name_keywords(n, args.lr_backbone_names)
                and p.requires_grad
            ],
            'lr':
            args.lr_backbone,
        }, {
            'params': [
                p for n, p in model_without_ddp.named_parameters()
                if match_name_keywords(n, args.lr_linear_proj_names)
                and p.requires_grad
            ],
            'lr':
            args.lr * args.lr_linear_proj_mult,
        }]
        return param_dicts

    if param_dict_type == 'large_wd':
        param_dicts = [{
            'params': [
                p for n, p in model_without_ddp.named_parameters()
                if not match_name_keywords(n, ['backbone'])
                and not match_name_keywords(n, ['norm', 'bias'])
                and p.requires_grad
            ],
        }, {
            'params': [
                p for n, p in model_without_ddp.named_parameters()
                if match_name_keywords(n, ['backbone']) and
                match_name_keywords(n, ['norm', 'bias']) and p.requires_grad
            ],
            'lr':
            args.lr_backbone,
            'weight_decay':
            0.0,
        }, {
            'params': [
                p for n, p in model_without_ddp.named_parameters()
                if match_name_keywords(n, ['backbone'])
                and not match_name_keywords(n, ['norm', 'bias'])
                and p.requires_grad
            ],
            'lr':
            args.lr_backbone,
            'weight_decay':
            args.weight_decay,
        }, {
            'params': [
                p for n, p in model_without_ddp.named_parameters()
                if not match_name_keywords(n, ['backbone']) and
                match_name_keywords(n, ['norm', 'bias']) and p.requires_grad
            ],
            'lr':
            args.lr,
            'weight_decay':
            0.0,
        }]

    return param_dicts