google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""Model utils."""
17
import torch
18

19

20
def apply_weight_decay(net, weight_decay):
21
  """Apply weight decay."""
22
  if weight_decay == 0:
23
    return
24
  for _, param in net.named_parameters():
25
    if param.grad is None:
26
      continue
27
    param.grad = param.grad.add(param, alpha=weight_decay)
28

29

30
def load_model(net, kwargs):
31
  """Load pretrained model."""
32
  pretrained_path = kwargs.get('pretrained_path', False)
33
  assert pretrained_path, 'Could not find pretrained_path!'
34
  print(f'Loading model from {pretrained_path}')
35
  state_dict = torch.load(pretrained_path)['state_dict']
36
  net.load_state_dict(state_dict)
37
  return net
38