7412047575

Committed 04 Jan 2024 04:18PM UTC coverage: 83.826% (-0.4%) from 84.211%

Build # 7412047575

Build Type

push

github

Committed by

hahnec

Commit Message

fix(ver): consider torch version for function_call import

Run Details

120 of 159 branches covered (0.0%)

Branch coverage included in aggregate %.

4 of 9 new or added lines in 1 file covered. (44.44%)

15 existing lines in 1 file now uncovered.

362 of 416 relevant lines covered (87.02%)

2.61 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

66.07

/torchimize/optimizer/gna_opt.py

import torch
from torch import Tensor
import torch.nn as nn
from torch.optim.optimizer import Optimizer
from typing import List

from distutils.version import LooseVersion
if LooseVersion(torch.__version__) > LooseVersion('2.0.0'):
    from torch.func import functional_call
else:
    try:
        from torch.nn.utils.stateless import functional_call
    except ImportError:
        from torch.nn.utils._stateless import functional_call



class GNA(Optimizer):
    r"""Implements Gauss-Newton.
    """

    def __init__(self, params, lr: float, model: nn.Module, hessian_approx: bool = True):

        if lr is not None and lr < 0.0:
            raise ValueError("Invalid learning rate: {}".format(lr))

        defaults = dict(lr=lr)

        super(GNA, self).__init__(params, defaults)

        self.hessian_approx = hessian_approx

        self._model = model
        self._params = self.param_groups[0]['params']
        self._j_list = []
        self._h_list = []

    def __setstate__(self, state):
        super(GNA, self).__setstate__(state)
        for group in self.param_groups:
            group.setdefault('nesterov', False)
            group.setdefault('amsgrad', False)

    @torch.no_grad()
    def step(self, x: torch.Tensor, closure=None):
        """Performs a single optimization step.

        Args:
            x: Current data batch, which is needed to compute 2nd order derivatives
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            params_with_grad = []
            d_p_list = []
            lr = group['lr']

            for p in group['params']:
                if p.grad is not None:
                    params_with_grad.append(p)
                    d_p_list.append(p.grad)

            parameters = dict(self._model.named_parameters())
            keys, values = zip(*parameters.items())

            self._h_list = []
            if self.hessian_approx:
                # vectorized jacobian (https://github.com/pytorch/pytorch/issues/49171)
                def func(*params: torch.Tensor):
                    out = functional_call(self._model, {n: p for n, p in zip(keys, params)}, x)
                    return out
                self._j_list: tuple[torch.Tensor] = torch.autograd.functional.jacobian(func, values, create_graph=False)    # NxCxBxCxHxW
                # create hessian approximation
                for i, j in enumerate(self._j_list):
                    j = j.flatten(end_dim=len(self._j_list[i].shape)-len(d_p_list[i].shape)-1).flatten(start_dim=1)  # (NC)x(BCHW)
                    try:
                        h = j.T.matmul(j)
                    except RuntimeError:
                        h = None
                    self._h_list.append(h)
                
            else:
                # vectorized hessian (https://github.com/pytorch/pytorch/issues/49171)
                def func(*params: torch.Tensor):
                    out: torch.Tensor = functional_call(self._model, {n: p for n, p in zip(keys, params)}, x)
                    return out.square().sum()
                self._h_list: tuple[torch.Tensor] = torch.autograd.functional.hessian(func, tuple(self._model.parameters()), create_graph=False)
                self._h_list = [self._h_list[i][i] for i in range(len(self._h_list))] # filter j-th element
                self._h_list = [h.flatten(end_dim=len(self._h_list[i].shape)-len(d_p_list[i].shape)-1).flatten(start_dim=1) for i, h in enumerate(self._h_list)] # (NC)x(BCHW)

            self.gna_update(
                params_with_grad,
                d_p_list,
                lr=lr,
            )

        return loss

    def gna_update(
            self,
            params: List[Tensor],
            d_p_list: List[Tensor],
            lr: float,
                 ):
        r"""Functional API that performs Gauss-Newton algorithm computation.

        """

        assert len(d_p_list) == len(self._h_list), 'Layer number mismatch'

        # e.g. 4 list elements in raw_train_fit.py for hidden and predict weights and biases
        for i, param in enumerate(params):

            d_p = d_p_list[i]
            h = self._h_list[i]
            if h is None:
                # fall back to stochastic gradient descent
                param.add_(d_p, alpha=-lr)
                break
            # prevent zeros along Hessian diagonal
            diag_vec = h.diagonal() + torch.finfo(h.dtype).eps * 1
            h.as_strided([h.size(0)], [h.size(0) + 1]).copy_(diag_vec)
            h_i = h.pinverse()
            if h_i.shape[-1] == d_p.flatten().shape[0]:
                d2_p = h_i.matmul(d_p.flatten()).reshape(d_p_list[i].shape)
                param.add_(d2_p, alpha=-lr)
            else:
                raise Exception('Tensor dimension mismatch')

1	import torch	3✔
2	from torch import Tensor	3✔
3	import torch.nn as nn	3✔
4	from torch.optim.optimizer import Optimizer	3✔
5	from typing import List	3✔
6
7	from distutils.version import LooseVersion	3✔
8	if LooseVersion(torch.__version__) > LooseVersion('2.0.0'):	3!
9	from torch.func import functional_call	3✔
10	else:
NEW 11	try:	×
NEW 12	from torch.nn.utils.stateless import functional_call	×
NEW 13	except ImportError:	×
NEW 14	from torch.nn.utils._stateless import functional_call	×
15
16
17
18	class GNA(Optimizer):	3✔
19	r"""Implements Gauss-Newton.
20	"""
21
22	def __init__(self, params, lr: float, model: nn.Module, hessian_approx: bool = True):	3✔
23
24	if lr is not None and lr < 0.0:	3!
UNCOV 25	raise ValueError("Invalid learning rate: {}".format(lr))	×
26
27	defaults = dict(lr=lr)	3✔
28
29	super(GNA, self).__init__(params, defaults)	3✔
30
31	self.hessian_approx = hessian_approx	3✔
32
33	self._model = model	3✔
34	self._params = self.param_groups[0]['params']	3✔
35	self._j_list = []	3✔
36	self._h_list = []	3✔
37
38	def __setstate__(self, state):	3✔
UNCOV 39	super(GNA, self).__setstate__(state)	×
UNCOV 40	for group in self.param_groups:	×
UNCOV 41	group.setdefault('nesterov', False)	×
UNCOV 42	group.setdefault('amsgrad', False)	×
43
44	@torch.no_grad()	3✔
45	def step(self, x: torch.Tensor, closure=None):	3✔
46	"""Performs a single optimization step.
47
48	Args:
49	x: Current data batch, which is needed to compute 2nd order derivatives
50	closure (callable, optional): A closure that reevaluates the model
51	and returns the loss.
52	"""
53	loss = None	3✔
54	if closure is not None:	3!
UNCOV 55	with torch.enable_grad():	×
UNCOV 56	loss = closure()	×
57
58	for group in self.param_groups:	3✔
59	params_with_grad = []	3✔
60	d_p_list = []	3✔
61	lr = group['lr']	3✔
62
63	for p in group['params']:	3✔
64	if p.grad is not None:	3!
65	params_with_grad.append(p)	3✔
66	d_p_list.append(p.grad)	3✔
67
68	parameters = dict(self._model.named_parameters())	3✔
69	keys, values = zip(*parameters.items())	3✔
70
71	self._h_list = []	3✔
72	if self.hessian_approx:	3!
73	# vectorized jacobian (https://github.com/pytorch/pytorch/issues/49171)
UNCOV 74	def func(*params: torch.Tensor):	×
NEW 75	out = functional_call(self._model, {n: p for n, p in zip(keys, params)}, x)	×
UNCOV 76	return out	×
UNCOV 77	self._j_list: tuple[torch.Tensor] = torch.autograd.functional.jacobian(func, values, create_graph=False) # NxCxBxCxHxW	×
78	# create hessian approximation
UNCOV 79	for i, j in enumerate(self._j_list):	×
80	j = j.flatten(end_dim=len(self._j_list[i].shape)-len(d_p_list[i].shape)-1).flatten(start_dim=1) # (NC)x(BCHW)	×
81	try:	×
82	h = j.T.matmul(j)	×
83	except RuntimeError:	×
UNCOV 84	h = None	×
85	self._h_list.append(h)	×
86
87	else:
88	# vectorized hessian (https://github.com/pytorch/pytorch/issues/49171)
89	def func(*params: torch.Tensor):	3✔
90	out: torch.Tensor = functional_call(self._model, {n: p for n, p in zip(keys, params)}, x)	3✔
91	return out.square().sum()	3✔
92	self._h_list: tuple[torch.Tensor] = torch.autograd.functional.hessian(func, tuple(self._model.parameters()), create_graph=False)	3✔
93	self._h_list = [self._h_list[i][i] for i in range(len(self._h_list))] # filter j-th element	3✔
94	self._h_list = [h.flatten(end_dim=len(self._h_list[i].shape)-len(d_p_list[i].shape)-1).flatten(start_dim=1) for i, h in enumerate(self._h_list)] # (NC)x(BCHW)	3✔
95
96	self.gna_update(	3✔
97	params_with_grad,
98	d_p_list,
99	lr=lr,
100	)
101
102	return loss	3✔
103
104	def gna_update(	3✔
105	self,
106	params: List[Tensor],
107	d_p_list: List[Tensor],
108	lr: float,
109	):
110	r"""Functional API that performs Gauss-Newton algorithm computation.
111
112	"""
113
114	assert len(d_p_list) == len(self._h_list), 'Layer number mismatch'	3✔
115
116	# e.g. 4 list elements in raw_train_fit.py for hidden and predict weights and biases
117	for i, param in enumerate(params):	3✔
118
119	d_p = d_p_list[i]	3✔
120	h = self._h_list[i]	3✔
121	if h is None:	3!
122	# fall back to stochastic gradient descent
UNCOV 123	param.add_(d_p, alpha=-lr)	×
UNCOV 124	break	×
125	# prevent zeros along Hessian diagonal
126	diag_vec = h.diagonal() + torch.finfo(h.dtype).eps * 1	3✔
127	h.as_strided([h.size(0)], [h.size(0) + 1]).copy_(diag_vec)	3✔
128	h_i = h.pinverse()	3✔
129	if h_i.shape[-1] == d_p.flatten().shape[0]:	3!
130	d2_p = h_i.matmul(d_p.flatten()).reshape(d_p_list[i].shape)	3✔
131	param.add_(d2_p, alpha=-lr)	3✔
132	else:
UNCOV 133	raise Exception('Tensor dimension mismatch')	×

hahnec / torchimize / 7412047575

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous