-
Notifications
You must be signed in to change notification settings - Fork 512
/
base_layer.py
63 lines (50 loc) · 1.94 KB
/
base_layer.py
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
#
# For licensing see accompanying LICENSE file.
# Copyright (C) 2024 Apple Inc. All Rights Reserved.
#
import argparse
from typing import Any, Dict, List, Optional, Tuple
from torch import nn
from corenet.modeling.misc.common import parameter_list
class BaseLayer(nn.Module):
"""
Base class for neural network layers. Subclass must implement `forward` function.
"""
def __init__(self, *args, **kwargs) -> None:
super().__init__()
@classmethod
def add_arguments(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
"""Add layer specific arguments"""
return parser
def get_trainable_parameters(
self,
weight_decay: Optional[float] = 0.0,
no_decay_bn_filter_bias: Optional[bool] = False,
*args,
**kwargs
) -> Tuple[List[Dict], List[float]]:
"""
Get parameters for training along with the learning rate.
Args:
weight_decay: weight decay
no_decay_bn_filter_bias: Do not decay BN and biases. Defaults to False.
Returns:
Returns a tuple of length 2. The first entry is a list of dictionary with three keys
(params, weight_decay, param_names). The second entry is a list of floats containing
learning rate for each parameter.
Note:
Learning rate multiplier is set to 1.0 here as it is handled inside the Central Model.
"""
param_list = parameter_list(
named_parameters=self.named_parameters,
weight_decay=weight_decay,
no_decay_bn_filter_bias=no_decay_bn_filter_bias,
*args,
**kwargs
)
return param_list, [1.0] * len(param_list)
def forward(self, *args, **kwargs) -> Any:
"""Forward function."""
raise NotImplementedError("Sub-classes should implement forward method")
def __repr__(self):
return "{}".format(self.__class__.__name__)