Coverage for zanj/torchutil.py: 97%

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

23 from zanj.consts import _FORMAT_KEY, safe_getsource, string_as_lines

24

25 from zanj import ZANJ, register_loader_handler

26 from zanj.loading import LoaderHandler, load_item_recursive

27

28# pylint: disable=protected-access

29

30 KWArgs = Any

31

32

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

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

35

36 - only counting shared parameters once

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

38

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

40 """

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

42 if only_trainable:

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

44

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

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

47 )

48

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

50

51

52 def get_module_device(

53 m: torch.nn.Module,

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

55 """get the current devices"""

56

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

58

59 if len(devs) == 0:

60 return False, devs

61

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

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

64

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

66 return True, dev_uni

67 else:

68 return False, devs

69

70

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

72

73

74 class ConfiguredModel(

75 torch.nn.Module,

76 typing.Generic[T_config],

77 metaclass=abc.ABCMeta,

78):

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

80

81 ```python

82 @set_config_class(YourConfig)

83 class YourModule(ConfiguredModel[YourConfig]):

84 def __init__(self, cfg: YourConfig):

85 super().__init__(cfg)

86 ```

87

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

89 `super().__init__(zanj_model_config)`

90

91 If you are inheriting from another class + ConfiguredModel,

92 ConfiguredModel must be the first class in the inheritance list

93 """

94

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

96 _config_class: type | None = None

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

98

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

100 super().__init__(**kwargs)

101 if self.zanj_config_class is None:

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

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

104 raise TypeError(

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

106 )

107

108 self.zanj_model_config: T_config = zanj_model_config

109 self.training_records: dict | None = None

110

111 def serialize(

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

113 ) -> dict[str, Any]:

114 if zanj is None:

115 zanj = ZANJ()

116 obj = dict(

117 zanj_model_config=self.zanj_model_config.serialize(),

118 meta=dict(

119 class_name=self.__class__.__name__,

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

121 class_source=safe_getsource(self.__class__),

122 module_name=self.__class__.__module__,

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

124 num_params=num_params(self),

125 as_str=string_as_lines(str(self)),

126 ),

127 training_records=self.training_records,

128 state_dict=self.state_dict(),

129 __muutils_format__=self.__class__.__name__,

130 )

131 return obj

132

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

134 if zanj is None:

135 zanj = ZANJ()

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

137

138 def _load_state_dict_wrapper(

139 self,

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

141 **kwargs,

142 ):

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

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

145 return self.load_state_dict(state_dict)

146

147 @classmethod

148 def load(

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

150 ) -> "ConfiguredModel":

151 """load a model from a serialized object"""

152

153 if zanj is None:

154 zanj = ZANJ()

155

156 # get the config

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

158

159 # get the training records

160 training_records: typing.Any = load_item_recursive(

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

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

163 zanj,

164 )

165

166 # initialize the model

167 model: "ConfiguredModel" = cls(zanj_model_config)

168

169 # load the state dict

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

171 obj["state_dict"],

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

173 zanj,

174 )

175

176 model._load_state_dict_wrapper(

177 tensored_state_dict,

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

179 )

180

181 # set the training records

182 model.training_records = training_records

183

184 return model

185

186 @classmethod

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

188 """read a model from a file"""

189 if zanj is None:

190 zanj = ZANJ()

191

192 mdl: ConfiguredModel = zanj.read(file_path)

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

194 return mdl

195

196 @classmethod

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

198 """read a model from a file"""

199 warnings.warn(

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

201 )

202 return cls.read(file_path, zanj)

203

204 @classmethod

205 def get_handler(cls) -> LoaderHandler:

206 cls_name: str = str(cls.__name__)

207 return LoaderHandler(

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

209 isinstance(json_item, dict)

210 and _FORMAT_KEY in json_item

211 and json_item[_FORMAT_KEY].startswith(cls_name)

212 ),

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

214 uid=cls_name,

215 source_pckg=cls.__module__,

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

217 )

218

219 def num_params(self) -> int:

220 return num_params(self)

221

222

223 def set_config_class(

224 config_class: Type[SerializableDataclass],

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

226 if not issubclass(config_class, SerializableDataclass):

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

228

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

230 # set the config class

231 cls._config_class = config_class

232

233 # register the handlers

234 register_loader_handler(cls.get_handler())

235

236 # return the new class

237 return cls

238

239 return wrapper

240

241

242 class ConfigMismatchException(ValueError):

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

244 super().__init__(msg)

245 self.diff = diff

246

247 def __str__(self):

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

249

250

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

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

253

254 Raises:

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

256 """

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

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

259 "model_a must have a zanj_model_config"

260 )

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

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

263 "model_b must have a zanj_model_config"

264 )

265

266 cls_type: type = type(model_a.zanj_model_config)

267

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

269 raise ConfigMismatchException(

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

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

272 )

273

274

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

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

277 assert_model_cfg_equality(model_a, model_b)

278

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

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

281 assert model_a_sd_keys == model_b_sd_keys, (

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

283 )

284 keys_failed: list[str] = list()

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

286 v_b = model_b.state_dict()[k]

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

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

289 keys_failed.append(k)

290 print(f"failed {k}")

291 else:

292 print(f"passed {k}")

293 assert len(keys_failed) == 0, (

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

295 )