Coverage for zanj/torchutil.py: 90%

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 

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

23 

24from zanj import ZANJ, register_loader_handler 

25from zanj.loading import LoaderHandler, load_item_recursive 

26 

27# pylint: disable=protected-access 

28 

29KWArgs = Any 

30 

31 

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

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

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

71 

72 

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

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

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

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

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