Coverage for zanj / torchutil.py: 97%

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

2 

3note that this requires torch 

4""" 

5 

6from __future__ import annotations 

7 

8import abc 

9import typing 

10import warnings 

11from typing import Any, Type, TypeVar 

12 

13try: 

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

15except 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 

20from muutils.json_serialize import SerializableDataclass 

21from muutils.json_serialize.json_serialize import ObjectPath 

22 

23from zanj.consts import _FORMAT_KEY, safe_getsource, string_as_lines 

24 

25from zanj import ZANJ, register_loader_handler 

26from zanj.loading import LoaderHandler, load_item_recursive 

27 

28# pylint: disable=protected-access 

29 

30KWArgs = Any 

31 

32 

33def 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 

52def 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 

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

72 

73 

74class 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 

223def 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 

242class 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 

251def 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 

275def 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 )