5600673849

################################################################################
# Copyright (c) 2021 ContinualAI.                                              #
# Copyrights licensed under the MIT License.                                   #
# See the accompanying LICENSE file for terms.                                 #
#                                                                              #
# Date: 30-12-2020                                                             #
# Author(s): Diganta Misra                                                     #
# E-mail: contact@continualai.org                                              #
# Website: www.continualai.org                                                 #
################################################################################

import copy
import io
from typing import TYPE_CHECKING, Optional

from torch import Tensor
import torch

from avalanche.evaluation import PluginMetric
from avalanche.evaluation.metric_results import MetricValue, MetricResult
from avalanche.evaluation.metric_utils import get_metric_name

if TYPE_CHECKING:
    from avalanche.training.templates import SupervisedTemplate


class WeightCheckpoint(PluginMetric[Tensor]):
    """
    The WeightCheckpoint Metric.

    Instances of this metric keeps the weight checkpoint tensor of the
    model at each experience.

    Each time `result` is called, this metric emits the latest experience's
    weight checkpoint tensor since the last `reset`.

    The reset method will bring the metric to its initial state. By default
    this metric in its initial state will return None.
    """

    def __init__(self):
        """
        Creates an instance of the WeightCheckpoint Metric.

        By default this metric in its initial state will return None.
        The metric can be updated by using the `update` method
        while the current experience's weight checkpoint tensor can be
        retrieved using the `result` method.
        """
        super().__init__()
        self.weights: Optional[bytes] = None

    def update(self, weights: bytes):
        """
        Update the weight checkpoint at the current experience.

        :param weights: the weight tensor at current experience
        :return: None.
        """
        self.weights = weights

    def result(self) -> Optional[bytes]:
        """
        Retrieves the weight checkpoint at the current experience.

        :return: The weight checkpoint as a tensor.
        """
        return self.weights

    def reset(self) -> None:
        """
        Resets the metric.

        :return: None.
        """
        self.weights = None

    def _package_result(self, strategy) -> "MetricResult":
        weights = self.result()
        if weights is None:
            return None

        metric_name = get_metric_name(
            self, strategy, add_experience=True, add_task=False
        )
        return [
            MetricValue(self, metric_name, weights, strategy.clock.train_iterations)
        ]

    def after_training_exp(self, strategy: "SupervisedTemplate") -> "MetricResult":
        buff = io.BytesIO()
        model_params = copy.deepcopy(strategy.model).to("cpu")
        torch.save(model_params, buff)
        buff.seek(0)
        self.update(buff.read())

        return self._package_result(strategy)

    def __str__(self):
        return "WeightCheckpoint"


__all__ = ["WeightCheckpoint"]

1	################################################################################
2	# Copyright (c) 2021 ContinualAI. #
3	# Copyrights licensed under the MIT License. #
4	# See the accompanying LICENSE file for terms. #
5	# #
6	# Date: 30-12-2020 #
7	# Author(s): Diganta Misra #
8	# E-mail: contact@continualai.org #
9	# Website: www.continualai.org #
10	################################################################################
11
12	import copy	4✔
13	import io	4✔
14	from typing import TYPE_CHECKING, Optional	4✔
15
16	from torch import Tensor	4✔
17	import torch	4✔
18
19	from avalanche.evaluation import PluginMetric	4✔
20	from avalanche.evaluation.metric_results import MetricValue, MetricResult	4✔
21	from avalanche.evaluation.metric_utils import get_metric_name	4✔
22
23	if TYPE_CHECKING:	4✔
24	from avalanche.training.templates import SupervisedTemplate	×
25
26
27	class WeightCheckpoint(PluginMetric[Tensor]):	4✔
28	"""	4✔
29	The WeightCheckpoint Metric.
30
31	Instances of this metric keeps the weight checkpoint tensor of the
32	model at each experience.
33
34	Each time `result` is called, this metric emits the latest experience's
35	weight checkpoint tensor since the last `reset`.
36
37	The reset method will bring the metric to its initial state. By default
38	this metric in its initial state will return None.
39	"""
40
41	def __init__(self):	4✔
42	"""
43	Creates an instance of the WeightCheckpoint Metric.
44
45	By default this metric in its initial state will return None.
46	The metric can be updated by using the `update` method
47	while the current experience's weight checkpoint tensor can be
48	retrieved using the `result` method.
49	"""
50	super().__init__()	×
NEW 51	self.weights: Optional[bytes] = None	×
52
53	def update(self, weights: bytes):	4✔
54	"""
55	Update the weight checkpoint at the current experience.
56
57	:param weights: the weight tensor at current experience
58	:return: None.
59	"""
60	self.weights = weights	×
61
62	def result(self) -> Optional[bytes]:	4✔
63	"""
64	Retrieves the weight checkpoint at the current experience.
65
66	:return: The weight checkpoint as a tensor.
67	"""
68	return self.weights	×
69
70	def reset(self) -> None:	4✔
71	"""
72	Resets the metric.
73
74	:return: None.
75	"""
76	self.weights = None	×
77
78	def _package_result(self, strategy) -> "MetricResult":	4✔
79	weights = self.result()	×
NEW 80	if weights is None:	×
NEW 81	return None	×
82
UNCOV 83	metric_name = get_metric_name(	×
84	self, strategy, add_experience=True, add_task=False
85	)
86	return [	×
87	MetricValue(self, metric_name, weights, strategy.clock.train_iterations)
88	]
89
90	def after_training_exp(self, strategy: "SupervisedTemplate") -> "MetricResult":	4✔
NEW 91	buff = io.BytesIO()	×
NEW 92	model_params = copy.deepcopy(strategy.model).to("cpu")	×
NEW 93	torch.save(model_params, buff)	×
NEW 94	buff.seek(0)	×
NEW 95	self.update(buff.read())	×
96
NEW 97	return self._package_result(strategy)	×
98
99	def __str__(self):	4✔
100	return "WeightCheckpoint"	×
101
102
103	__all__ = ["WeightCheckpoint"]	4✔

ContinualAI / avalanche / 5600673849

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