Coverage for zanj/torchutil.py: 90%

1"""torch utilities for zanj -- in particular the `ConfiguredModel` base class

2

3 note that this requires torch

4"""

5

6 from __future__ import annotations

7

8 import abc

9 import typing

10 import warnings

11 from typing import Any, Type, TypeVar

12

13 try:

14 import torch # type: ignore[import-not-found]

15 except ImportError as e:

16 raise ImportError(

17 "torch is required for zanj.torchutil, please install it with `pip install torch` or `pip install zanj[torch]`"

18 ) from e

19

20 from muutils.json_serialize import SerializableDataclass

21 from muutils.json_serialize.json_serialize import ObjectPath

22 from muutils.json_serialize.util import safe_getsource, string_as_lines, _FORMAT_KEY

23

24 from zanj import ZANJ, register_loader_handler

25 from zanj.loading import LoaderHandler, load_item_recursive

26

27# pylint: disable=protected-access

28

29 KWArgs = Any

30

31

32 def num_params(m: torch.nn.Module, only_trainable: bool = True):

33 """return total number of parameters in a model

34

35 - only counting shared parameters once

36 - if `only_trainable` is False, will include parameters with `requires_grad = False`

37

38 https://stackoverflow.com/questions/49201236/check-the-total-number-of-parameters-in-a-pytorch-model

39 """

40 parameters: list[torch.nn.Parameter] = list(m.parameters())

41 if only_trainable:

42 parameters = [p for p in parameters if p.requires_grad]

43

44 unique: list[torch.nn.Parameter] = list(

45 {p.data_ptr(): p for p in parameters}.values()

46 )

47

48 return sum(p.numel() for p in unique)

49

50

51 def get_module_device(

52 m: torch.nn.Module,

53) -> tuple[bool, torch.device | dict[str, torch.device]]:

54 """get the current devices"""

55

56 devs: dict[str, torch.device] = {name: p.device for name, p in m.named_parameters()}

57

58 if len(devs) == 0:

59 return False, devs

60

61 # check if all devices are the same by getting one device

62 dev_uni: torch.device = next(iter(devs.values()))

63

64 if all(dev == dev_uni for dev in devs.values()):

65 return True, dev_uni

66 else:

67 return False, devs

68

69

70 T_config = TypeVar("T_config", bound=SerializableDataclass)

71

72

73 class ConfiguredModel(

74 torch.nn.Module,

75 typing.Generic[T_config],

76 metaclass=abc.ABCMeta,

77):

78 """a model that has a configuration, for saving with ZANJ

79

80 ```python

81 @set_config_class(YourConfig)

82 class YourModule(ConfiguredModel[YourConfig]):

83 def __init__(self, cfg: YourConfig):

84 super().__init__(cfg)

85 ```

86

87 `__init__()` must initialize the model from a config object only, and call

88 `super().__init__(zanj_model_config)`

89

90 If you are inheriting from another class + ConfiguredModel,

91 ConfiguredModel must be the first class in the inheritance list

92 """

93

94 # dont set this directly, use `set_config_class()` decorator

95 _config_class: type | None = None

96 zanj_config_class = property(lambda self: type(self)._config_class)

97

98 def __init__(self, zanj_model_config: T_config, **kwargs):

99 super().__init__(**kwargs)

100 if self.zanj_config_class is None:

101 raise NotImplementedError("you need to set `config_class` for your model")

102 if not isinstance(zanj_model_config, self.zanj_config_class): # type: ignore

103 raise TypeError(

104 f"config must be an instance of {self.zanj_config_class = }, got {type(zanj_model_config) = }"

105 )

106

107 self.zanj_model_config: T_config = zanj_model_config

108 self.training_records: dict | None = None

109

110 def serialize(

111 self, path: ObjectPath = tuple(), zanj: ZANJ | None = None

112 ) -> dict[str, Any]:

113 if zanj is None:

114 zanj = ZANJ()

115 obj = dict(

116 zanj_model_config=self.zanj_model_config.serialize(),

117 meta=dict(

118 class_name=self.__class__.__name__,

119 class_doc=string_as_lines(self.__class__.__doc__),

120 class_source=safe_getsource(self.__class__),

121 module_name=self.__class__.__module__,

122 module_mro=[str(x) for x in self.__class__.__mro__],

123 num_params=num_params(self),

124 as_str=string_as_lines(str(self)),

125 ),

126 training_records=self.training_records,

127 state_dict=self.state_dict(),

128 __muutils_format__=self.__class__.__name__,

129 )

130 return obj

131

132 def save(self, file_path: str, zanj: ZANJ | None = None):

133 if zanj is None:

134 zanj = ZANJ()

135 zanj.save(self.serialize(), file_path)

136

137 def _load_state_dict_wrapper(

138 self,

139 state_dict: dict[str, torch.Tensor],

140 **kwargs,

141 ):

142 """wrapper for `load_state_dict()` in case you need to override it"""

143 assert len(kwargs) == 0, f"got unexpected kwargs: {kwargs}"

144 return self.load_state_dict(state_dict)

145

146 @classmethod

147 def load(

148 cls, obj: dict[str, Any], path: ObjectPath, zanj: ZANJ | None = None

149 ) -> "ConfiguredModel":

150 """load a model from a serialized object"""

151

152 if zanj is None:

153 zanj = ZANJ()

154

155 # get the config

156 zanj_model_config: T_config = cls._config_class.load(obj["zanj_model_config"]) # type: ignore

157

158 # get the training records

159 training_records: typing.Any = load_item_recursive(

160 obj.get("training_records", None),

161 tuple(path) + ("training_records",),

162 zanj,

163 )

164

165 # initialize the model

166 model: "ConfiguredModel" = cls(zanj_model_config)

167

168 # load the state dict

169 tensored_state_dict: dict[str, torch.Tensor] = load_item_recursive(

170 obj["state_dict"],

171 tuple(path) + ("state_dict",),

172 zanj,

173 )

174

175 model._load_state_dict_wrapper(

176 tensored_state_dict,

177 **zanj.custom_settings.get("_load_state_dict_wrapper", dict()),

178 )

179

180 # set the training records

181 model.training_records = training_records

182

183 return model

184

185 @classmethod

186 def read(cls, file_path: str, zanj: ZANJ | None = None) -> "ConfiguredModel":

187 """read a model from a file"""

188 if zanj is None:

189 zanj = ZANJ()

190

191 mdl: ConfiguredModel = zanj.read(file_path)

192 assert isinstance(mdl, cls), f"loaded object must be a {cls}, got {type(mdl)}"

193 return mdl

194

195 @classmethod

196 def load_file(cls, file_path: str, zanj: ZANJ | None = None) -> "ConfiguredModel":

197 """read a model from a file"""

198 warnings.warn(

199 "load_file() is deprecated, use read() instead", DeprecationWarning

200 )

201 return cls.read(file_path, zanj)

202

203 @classmethod

204 def get_handler(cls) -> LoaderHandler:

205 cls_name: str = str(cls.__name__)

206 return LoaderHandler(

207 check=lambda json_item, path=None, z=None: ( # type: ignore

208 isinstance(json_item, dict)

209 and _FORMAT_KEY in json_item

210 and json_item[_FORMAT_KEY].startswith(cls_name)

211 ),

212 load=lambda json_item, path=None, z=None: cls.load(json_item, path, z), # type: ignore

213 uid=cls_name,

214 source_pckg=cls.__module__,

215 desc=f"{cls.__module__} {cls_name} loader via zanj.torchutil.ConfiguredModel",

216 )

217

218 def num_params(self) -> int:

219 return num_params(self)

220

221

222 def set_config_class(

223 config_class: Type[SerializableDataclass],

224) -> typing.Callable[[Type[ConfiguredModel]], Type[ConfiguredModel]]:

225 if not issubclass(config_class, SerializableDataclass):

226 raise TypeError(f"{config_class} must be a subclass of SerializableDataclass")

227

228 def wrapper(cls: Type[ConfiguredModel]) -> Type[ConfiguredModel]:

229 # set the config class

230 cls._config_class = config_class

231

232 # register the handlers

233 register_loader_handler(cls.get_handler())

234

235 # return the new class

236 return cls

237

238 return wrapper

239

240

241 class ConfigMismatchException(ValueError):

242 def __init__(self, msg: str, diff):

243 super().__init__(msg)

244 self.diff = diff

245

246 def __str__(self):

247 return f"{super().__str__()}: {self.diff}"

248

249

250 def assert_model_cfg_equality(model_a: ConfiguredModel, model_b: ConfiguredModel):

251 """check both models are correct instances and have the same config

252

253 Raises:

254 ConfigMismatchException: if the configs don't match, e.diff will contain the diff

255 """

256 assert isinstance(model_a, ConfiguredModel), "model_a must be a ConfiguredModel"

257 assert isinstance(model_a.zanj_model_config, SerializableDataclass), (

258 "model_a must have a zanj_model_config"

259 )

260 assert isinstance(model_b, ConfiguredModel), "model_b must be a ConfiguredModel"

261 assert isinstance(model_b.zanj_model_config, SerializableDataclass), (

262 "model_b must have a zanj_model_config"

263 )

264

265 cls_type: type = type(model_a.zanj_model_config)

266

267 if not (model_a.zanj_model_config == model_b.zanj_model_config):

268 raise ConfigMismatchException(

269 f"configs of type {type(model_a.zanj_model_config)}, {type(model_b.zanj_model_config)} don't match",

270 diff=cls_type.diff(model_a.zanj_model_config, model_b.zanj_model_config), # type: ignore[attr-defined]

271 )

272

273

274 def assert_model_exact_equality(model_a: ConfiguredModel, model_b: ConfiguredModel):

275 """check the models are exactly equal, including state dict contents"""

276 assert_model_cfg_equality(model_a, model_b)

277

278 model_a_sd_keys: set[str] = set(model_a.state_dict().keys())

279 model_b_sd_keys: set[str] = set(model_b.state_dict().keys())

280 assert model_a_sd_keys == model_b_sd_keys, (

281 f"state dict keys don't match: {model_a_sd_keys - model_b_sd_keys} / {model_b_sd_keys - model_a_sd_keys}"

282 )

283 keys_failed: list[str] = list()

284 for k, v_a in model_a.state_dict().items():

285 v_b = model_b.state_dict()[k]

286 if not (v_a == v_b).all():

287 # if not torch.allclose(v, v_load):

288 keys_failed.append(k)

289 print(f"failed {k}")

290 else:

291 print(f"passed {k}")

292 assert len(keys_failed) == 0, (

293 f"{len(keys_failed)} / {len(model_a_sd_keys)} state dict elements don't match: {keys_failed}"

294 )