Spaces:
Build error
Build error
| """Dynamic Buffer Module.""" | |
| # Copyright (C) 2020 Intel Corporation | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, | |
| # software distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions | |
| # and limitations under the License. | |
| from abc import ABC | |
| from torch import Tensor, nn | |
| class DynamicBufferModule(ABC, nn.Module): | |
| """Torch module that allows loading variables from the state dict even in the case of shape mismatch.""" | |
| def get_tensor_attribute(self, attribute_name: str) -> Tensor: | |
| """Get attribute of the tensor given the name. | |
| Args: | |
| attribute_name (str): Name of the tensor | |
| Raises: | |
| ValueError: `attribute_name` is not a torch Tensor | |
| Returns: | |
| Tensor: Tensor attribute | |
| """ | |
| attribute = self.__getattr__(attribute_name) | |
| if isinstance(attribute, Tensor): | |
| return attribute | |
| raise ValueError(f"Attribute with name '{attribute_name}' is not a torch Tensor") | |
| def _load_from_state_dict(self, state_dict: dict, prefix: str, *args): | |
| """Resizes the local buffers to match those stored in the state dict. | |
| Overrides method from parent class. | |
| Args: | |
| state_dict (dict): State dictionary containing weights | |
| prefix (str): Prefix of the weight file. | |
| *args: | |
| """ | |
| persistent_buffers = {k: v for k, v in self._buffers.items() if k not in self._non_persistent_buffers_set} | |
| local_buffers = {k: v for k, v in persistent_buffers.items() if v is not None} | |
| for param in local_buffers.keys(): | |
| for key in state_dict.keys(): | |
| if key.startswith(prefix) and key[len(prefix) :].split(".")[0] == param: | |
| if not local_buffers[param].shape == state_dict[key].shape: | |
| attribute = self.get_tensor_attribute(param) | |
| attribute.resize_(state_dict[key].shape) | |
| super()._load_from_state_dict(state_dict, prefix, *args) | |