-
Notifications
You must be signed in to change notification settings - Fork 1.7k
Expand file tree
/
Copy pathcheckpoints.py
More file actions
1388 lines (1214 loc) · 51.4 KB
/
checkpoints.py
File metadata and controls
1388 lines (1214 loc) · 51.4 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
"""This module implements a checkpoint saver and loader.
A checkpoint in an experiment usually needs to save the state of many different
things: the model parameters, optimizer parameters, what epoch is this, etc.
The save format for a checkpoint is a directory, where each of these separate
saveable things gets its own file. Additionally, a special file holds meta
information about the checkpoint (by default just time of creation, but you
can specify anything else you may wish, e.g. validation loss).
The interface for the checkpoint system requires you to specify what things to
save. This approach is flexible and agnostic of how your experiment is actually
run.
The interface requires you to specify names for each thing to save. This name
is used to give the right parameter file to the right object when recovering.
Default saving and loading methods are only added for torch.nn.Modules (and
their subclasses), and torch.optim.Optimizers. If those methods do not work for
your object, you can specify your own saving and/or loading methods, either for
a particular instance or a for a class.
Example
-------
>>> # Toy example Module:
>>> class Recoverable(torch.nn.Module):
... def __init__(self, param):
... super().__init__()
... self.param = torch.nn.Parameter(torch.tensor([param]))
...
... def forward(self, x):
... return x * self.param
>>> model = Recoverable(1.0)
>>> tempdir = getfixture("tmpdir")
>>> # In simple cases, the module aims to have a terse syntax,
>>> # consisting of three steps.
>>> # 1. Specifying where to save checkpoints and what is included in a
>>> # checkpoint:
>>> checkpointer = Checkpointer(tempdir, {"network": model})
>>> # 2. Recover from the latest checkpoint, if one is found:
>>> checkpointer.recover_if_possible()
>>> # Run your experiment:
>>> data = [(0.1, 0.9), (0.3, 0.8)]
>>> for example, target in data:
... loss = (model(example) - target) ** 2
... # 3. Save checkpoints, and keep by default just one, the newest:
... ckpt = checkpointer.save_and_keep_only()
Authors
* Aku Rouhe 2020
* Adel Moumen 2024
"""
import collections
import collections.abc
import inspect
import logging
import os
import pathlib
import shutil
import time
import warnings
from typing import Dict
import torch
import yaml
from packaging import version
import speechbrain.utils._workarounds as __wa
from speechbrain.utils.distributed import (
ddp_barrier,
ddp_broadcast,
if_main_process,
main_process_only,
once_per_node,
)
from speechbrain.utils.logger import get_logger
logger = get_logger(__name__)
CKPT_PREFIX = "CKPT"
METAFNAME = f"{CKPT_PREFIX}.yaml" # Important that this is not .ckpt
PARAMFILE_EXT = ".ckpt" # ...because these files will be
# some keys have been renamed in the new version of the code
KEYS_MAPPING: Dict[str, str] = {
".mutihead_attn": ".multihead_attn", # see PR #2489
".convs_intermedite": ".convs_intermediate", # fix for PostNet blame #2463
}
def map_old_state_dict_weights(
state_dict: Dict[str, torch.Tensor], mapping: Dict[str, str]
) -> Dict[str, torch.Tensor]:
"""
Maps the keys in the old state dictionary according to the provided mapping.
NOTE: This function will remap all state_dict keys that contain the old key.
For instance, if the state_dict is {'model.encoder.layer.0.atn.self.query.weight': ...}
and the mapping is {'.atn': '.attn'}, the resulting state_dict will be
{'model.encoder.layer.0.attn.self.query.weight': ...}.
Since this effectively works as a mass substring replacement, partial key
matches (e.g. in the middle of one layer name) will also work, so be
careful to avoid false positives.
Parameters
----------
state_dict : dict
The old state dictionary to be mapped.
mapping : dict
A dictionary specifying the mapping between old and new keys.
Returns
-------
dict
The modified state dictionary with mapped keys.
"""
for replacement_old, replacement_new in mapping.items():
for old_key in list(state_dict.keys()):
if replacement_old in old_key:
new_key = old_key.replace(replacement_old, replacement_new)
state_dict[new_key] = state_dict.pop(old_key)
logger.info(
"Due to replacement compatibility rule '%s'->'%s', renamed "
"`state_dict['%s']`->`state_dict['%s']`",
replacement_old,
replacement_new,
old_key,
new_key,
)
return state_dict
def hook_on_loading_state_dict_checkpoint(
state_dict: Dict[str, torch.Tensor],
) -> Dict[str, torch.Tensor]:
"""Hook to be called when loading a state_dict checkpoint.
This hook is called when loading a state_dict checkpoint. It can be used
to modify the state_dict before it is loaded into the model.
By default, this hook will map the old state_dict keys to the new ones.
Arguments
---------
state_dict : dict
The state_dict to be loaded.
Returns
-------
dict
The modified state_dict.
"""
altered_state_dict = map_old_state_dict_weights(state_dict, KEYS_MAPPING)
return altered_state_dict
def torch_recovery(obj, path, end_of_epoch):
"""Loads a torch.nn.Module state_dict from the given path instantly.
This can be made the default for torch.nn.Modules with:
>>> DEFAULT_LOAD_HOOKS[torch.nn.Module] = torch_recovery
Arguments
---------
obj : torch.nn.Module
Instance for which to load the parameters.
path : str, pathlib.Path
Path where to load from.
end_of_epoch : bool
Whether the recovery comes from an end of epoch checkpoint.
"""
del end_of_epoch # Unused
device = "cpu"
state_dict = torch_patched_state_dict_load(path, device)
try:
obj.load_state_dict(state_dict, strict=True)
except TypeError:
obj.load_state_dict(state_dict)
def torch_patched_state_dict_load(path, device="cpu"):
"""Loads a `state_dict` from the given path using :func:`torch.load` and
calls the SpeechBrain `state_dict` loading hooks, e.g. to apply key name
patching rules for compatibility.
The `state_dict` sees no further preprocessing and is not applied into a
model, see :func:`~torch_recovery` or :func:`~torch_parameter_transfer`.
Arguments
---------
path : str, pathlib.Path
Path where to load from.
device : str
Device where the loaded `state_dict` tensors should reside. This is
forwarded to :func:`torch.load`; see its documentation for details.
Returns
-------
The loaded state dict.
"""
state_dict = torch.load(path, map_location=device)
state_dict = hook_on_loading_state_dict_checkpoint(state_dict)
return state_dict
@main_process_only
def torch_save(obj, path):
"""Saves the obj's parameters to path.
Default save hook for torch.nn.Modules
For saving torch.nn.Module state_dicts.
Arguments
---------
obj : torch.nn.Module
Instance to save.
path : str, pathlib.Path
Path where to save to.
"""
state_dict = obj.state_dict()
if not state_dict:
logger.warning(f"Saving an empty state_dict for {obj} in {path}.")
torch.save(state_dict, path)
@once_per_node
def torch_save_once_per_node(obj, path):
"""Copy of `torch_save` that is run once per node."""
state_dict = obj.state_dict()
if not state_dict:
logger.warning(f"Saving an empty state_dict for {obj} in {path}.")
torch.save(state_dict, path)
def torch_parameter_transfer(obj, path):
"""Non-strict Torch Module state_dict load.
Loads a set of parameters from path to obj. If obj has layers for which
parameters can't be found, only a warning is logged. Same thing
if the path has parameters for layers which don't find a counterpart
in obj.
Arguments
---------
obj : torch.nn.Module
Instance for which to load the parameters.
path : str
Path where to load from.
"""
device = "cpu"
state_dict = torch_patched_state_dict_load(path, device)
incompatible_keys = obj.load_state_dict(state_dict, strict=False)
for missing_key in incompatible_keys.missing_keys:
logger.warning(
f"During parameter transfer to {obj} loading from "
+ f"{path}, the transferred parameters did not have "
+ f"parameters for the key: {missing_key}"
)
for unexpected_key in incompatible_keys.unexpected_keys:
logger.warning(
f"During parameter transfer to {obj} loading from "
+ f"{path}, the object could not use the parameters loaded "
+ f"with the key: {unexpected_key}"
)
# These dicts are indexed by class and hold the default checkpoints methods
DEFAULT_LOAD_HOOKS = {
torch.nn.Module: torch_recovery,
torch.optim.Optimizer: torch_recovery,
torch.optim.lr_scheduler.ReduceLROnPlateau: torch_recovery,
}
DEFAULT_SAVE_HOOKS = {
torch.nn.Module: torch_save,
torch.optim.Optimizer: torch_save,
torch.optim.lr_scheduler.ReduceLROnPlateau: torch_save,
}
if version.parse(torch.__version__) < version.parse("2.0.0"):
DEFAULT_LOAD_HOOKS[torch.optim.lr_scheduler._LRScheduler] = torch_recovery
DEFAULT_SAVE_HOOKS[torch.optim.lr_scheduler._LRScheduler] = torch_save
else:
DEFAULT_LOAD_HOOKS[torch.optim.lr_scheduler.LRScheduler] = torch_recovery
DEFAULT_SAVE_HOOKS[torch.optim.lr_scheduler.LRScheduler] = torch_save
if version.parse(torch.__version__) < version.parse("2.4.0"):
DEFAULT_LOAD_HOOKS[torch.cuda.amp.grad_scaler.GradScaler] = torch_recovery
DEFAULT_SAVE_HOOKS[torch.cuda.amp.grad_scaler.GradScaler] = torch_save
else:
DEFAULT_LOAD_HOOKS[torch.amp.grad_scaler.GradScaler] = torch_recovery
DEFAULT_SAVE_HOOKS[torch.amp.grad_scaler.GradScaler] = torch_save
DEFAULT_TRANSFER_HOOKS = {
torch.nn.Module: torch_parameter_transfer,
}
# Add a transfer hook for sentencepiece if it is installed:
try:
import sentencepiece as spm
def _load_spm(obj, path):
obj.load(str(path)) # SentencePieceProcessor needs a string.
DEFAULT_TRANSFER_HOOKS[spm.SentencePieceProcessor] = _load_spm
del spm # Don't leave it here bare.
except ImportError:
# SentencePiece not loaded, fine!
pass
# Add workarounds:
DEFAULT_SAVE_HOOKS[torch.optim.lr_scheduler.CyclicLR] = __wa._cycliclrsaver
DEFAULT_LOAD_HOOKS[torch.optim.lr_scheduler.CyclicLR] = __wa._cycliclrloader
def convert_torch_save_hooks_to_once_per_node():
"""Update the save hooks to be run once per node. This should be called
if you are running on more than one node with separate filesystems."""
global DEFAULT_SAVE_HOOKS
for obj, hook in DEFAULT_SAVE_HOOKS.items():
if hook == torch_save:
DEFAULT_SAVE_HOOKS[obj] = torch_save_once_per_node
def mark_as_saver(method):
"""Method decorator which marks given method as the checkpoint saving hook.
See register_checkpoint_hooks for example.
Arguments
---------
method : callable
Method of the class to decorate. Must be callable with
signature (instance, path) using positional arguments. This is
satisfied by for example: def saver(self, path):
Returns
-------
The decorated method, marked as a checkpoint saver.
Note
----
This will not add the hook (not possible via a method decorator),
you must also decorate the class with @register_checkpoint_hooks
Only one method can be added as the hook.
"""
sig = inspect.signature(method)
try:
sig.bind(object(), pathlib.Path("testpath"))
except TypeError:
MSG = "Checkpoint saver must match signature (instance, path)"
raise TypeError(MSG)
method._speechbrain_saver = True
return method
def mark_as_loader(method):
"""Method decorator which marks given method as checkpoint loading hook.
Arguments
---------
method : callable
Method of the class to decorate. Must be callable with
signature (instance, path, end_of_epoch) using positional
arguments. This is satisfied by for example:
`def loader(self, path, end_of_epoch):`
Returns
-------
The decorated method, registered as a checkpoint loader.
Note
----
This will not add the hook (not possible via a method decorator),
you must also decorate the class with @register_checkpoint_hooks
Only one method can be added as the hook.
"""
sig = inspect.signature(method)
try:
sig.bind(object(), pathlib.Path("testpath"), True)
except TypeError:
MSG = "Checkpoint loader must have signature (self, path, end_of_epoch)"
raise TypeError(MSG)
method._speechbrain_loader = True
return method
def mark_as_transfer(method):
"""Method decorator which marks given method as a parameter transfer hook.
Arguments
---------
method : callable
Method of the class to decorate. Must be callable with
signature (instance, path) using positional
arguments. This is satisfied by for example:
`def loader(self, path):`
Returns
-------
The decorated method, registered as a transfer method.
Note
----
This will not add the hook (not possible via a method decorator),
you must also decorate the class with @register_checkpoint_hooks
Only one method can be added as the hook.
Note
----
The transfer hook is prioritized over the loader hook by the ``Pretrainer``
However, if no transfer hook is registered, the Pretrainer will use the
loader hook.
"""
sig = inspect.signature(method)
try:
sig.bind(object(), pathlib.Path("testpath"))
except TypeError:
MSG = "Transfer hook must have signature (self, path)"
raise TypeError(MSG)
method._speechbrain_transfer = True
return method
def register_checkpoint_hooks(cls, save_on_main_only=True):
"""Class decorator which registers the load, save and transfer hooks.
The hooks must have been marked with mark_as_loader and mark_as_saver,
and possibly mark_as_transfer.
Arguments
---------
cls : class
Class to decorate
save_on_main_only : bool
By default, the saver is only run on a single process. This argument
provides the option to run the saver on all processes, needed
for some savers where data is first gathered before saving.
Returns
-------
the decorated class with hooks registered
Example
-------
>>> @register_checkpoint_hooks
... class CustomRecoverable:
... def __init__(self, param):
... self.param = int(param)
...
... @mark_as_saver
... def save(self, path):
... with open(path, "w", encoding="utf-8") as fo:
... fo.write(str(self.param))
...
... @mark_as_loader
... def load(self, path, end_of_epoch):
... del end_of_epoch # Unused here
... with open(path, encoding="utf-8") as fi:
... self.param = int(fi.read())
"""
global DEFAULT_LOAD_HOOKS
global DEFAULT_SAVE_HOOKS
global DEFAULT_TRANSFER_HOOKS
for name, method in cls.__dict__.items():
if hasattr(method, "_speechbrain_saver"):
# If the save method is to be run on main only, wrap the method with
# main_process_only() which stops it from running on the other procs
if save_on_main_only:
DEFAULT_SAVE_HOOKS[cls] = main_process_only(method)
else:
DEFAULT_SAVE_HOOKS[cls] = method
logger.debug(f"Registered checkpoint save hook for {name}")
if hasattr(method, "_speechbrain_loader"):
DEFAULT_LOAD_HOOKS[cls] = method
logger.debug(f"Registered checkpoint load hook for {name}")
if hasattr(method, "_speechbrain_transfer"):
DEFAULT_TRANSFER_HOOKS[cls] = method
logger.debug(f"Registered parameter transfer hook for {name}")
return cls
def get_default_hook(obj, default_hooks):
"""Finds the default save/load hook to use with the given object.
Follows the Method Resolution Order, i.e., if no hook is registered for
the class of the object itself, also searches classes which the object
inherits from.
Arguments
---------
obj : instance
Instance of a class.
default_hooks : dict
Mapping from classes to (checkpointing hook) functions.
Returns
-------
The correct method or None if no method is registered.
Example
-------
>>> a = torch.nn.Module()
>>> get_default_hook(a, DEFAULT_SAVE_HOOKS) == torch_save
True
"""
mro = inspect.getmro(type(obj))
for cls in mro:
if cls in default_hooks:
return default_hooks[cls]
# If we got here, no hook found
return None
Checkpoint = collections.namedtuple(
"Checkpoint", ["path", "meta", "paramfiles"]
)
Checkpoint.__doc__ = """NamedTuple describing one saved checkpoint
To select a checkpoint to load from many checkpoint,
Checkpoints are first filtered and sorted based on this namedtuple.
Checkpointers put pathlib.Path in path and a dict in meta.
You can essentially add any info you want to meta when saving a checkpoint.
The only default key in meta is "unixtime".
Checkpoint.paramfiles is a dict from recoverable name to parameter filepath.
"""
# Creating a hash allows making checkpoint sets
Checkpoint.__hash__ = lambda self: hash(self.path)
def ckpt_recency(ckpt):
"""Recency as Checkpoint importance metric.
This function can also act as an example of how to make checkpoint
importance keyfuncs. This is a named function, but as you can see
it could be easily implemented as a lambda in a pinch.
"""
return ckpt.meta["unixtime"]
class Checkpointer:
"""Saves checkpoints and recovers from them.
Arguments
---------
checkpoints_dir : str, pathlib.Path
Path to directory where to save checkpoints.
recoverables : mapping, optional
Objects to to recover. They need a (unique) name: this is used
to connect the parameters in a checkpoint to the correct recoverable.
The name is also used in the filename of the
savefile for the objects parameters. These can also be added with
add_recoverable or add_recoverables or just modifying
checkpointer.recoverables directly.
custom_load_hooks : mapping, optional
A mapping from name [same as in recoverables] to function or method.
Sets a custom loading hook for a particular object. The
function/method must be callable with signature (instance, path)
using positional arguments. This is satisfied by for example:
`def loader(self, path)`.
custom_save_hooks : mapping, optional
Mapping from name [same as in recoverables] to function or method.
Sets a custom saving hook for a particular object. The
function/method must be callable with
signature (instance, path) using positional arguments. This is
satisfied by for example: def saver(self, path):
allow_partial_load : bool, optional
If True, allows loading a checkpoint where a savefile is not found
for every registered recoverable. In that case, only the found
savefiles are loaded. When False, loading such a save will raise
RuntimeError. (default: False)
Example
-------
>>> import torch
>>> # SETUP:
>>> tempdir = getfixture("tmpdir")
>>> class Recoverable(torch.nn.Module):
... def __init__(self, param):
... super().__init__()
... self.param = torch.nn.Parameter(torch.tensor([param]))
...
... def forward(self, x):
... return x * self.param
>>> recoverable = Recoverable(1.0)
>>> recoverables = {"recoverable": recoverable}
>>> # SETUP DONE.
>>> checkpointer = Checkpointer(tempdir, recoverables)
>>> first_ckpt = checkpointer.save_checkpoint()
>>> recoverable.param.data = torch.tensor([2.0])
>>> loaded_ckpt = checkpointer.recover_if_possible()
>>> # Parameter has been loaded:
>>> assert recoverable.param.data == torch.tensor([1.0])
>>> # With this call, by default, oldest checkpoints are deleted:
>>> checkpointer.save_and_keep_only()
>>> assert first_ckpt not in checkpointer.list_checkpoints()
"""
def __init__(
self,
checkpoints_dir,
recoverables=None,
custom_load_hooks=None,
custom_save_hooks=None,
allow_partial_load=False,
):
self.checkpoints_dir = pathlib.Path(checkpoints_dir)
os.makedirs(self.checkpoints_dir, exist_ok=True)
self.recoverables = {}
self.optional_recoverables = {}
if recoverables is not None:
self.add_recoverables(recoverables)
self.custom_load_hooks = {}
if custom_load_hooks is not None:
self.custom_load_hooks.update(custom_load_hooks)
self.custom_save_hooks = {}
if custom_save_hooks is not None:
self.custom_save_hooks.update(custom_save_hooks)
self.allow_partial_load = allow_partial_load
def add_recoverable(
self,
name,
obj,
custom_load_hook=None,
custom_save_hook=None,
optional_load=False,
):
"""Register a recoverable with possible custom hooks.
Arguments
---------
name : str
Unique name for recoverable. Used to map savefiles to objects.
obj : instance
The object to recover.
custom_load_hook : callable, optional
Called to load the object's savefile. The function/method must be
callable with signature (instance, path) using positional
arguments. This is satisfied by for example: def load(self, path):
custom_save_hook : callable, optional
Called to save the object's parameters. The function/method must
be callable with signature (instance, path) using positional
arguments. This is satisfied by for example: def saver(self, path):
optional_load : bool, optional
If True, allows for the optional loading of an object from a checkpoint.
If the checkpoint lacks the specified object, no error is raised.
This is particularly useful during transitions between different training
configurations, such as changing precision from floating point 32 to 16.
For example, suppose you have a training checkpoint that does not includes
a `scaler` object. If you intend to continue pre-training in floating point 16,
where the `scaler` object is needed, marking it as optional prevents loading errors.
Without marking it as optional, attempting to load the `scaler` object from a checkpoint
trained in floating point 32 would fail, as the `scaler` object is not present
in that checkpoint.
"""
self.recoverables[name] = obj
self.optional_recoverables[name] = optional_load
if custom_load_hook is not None:
self.custom_load_hooks[name] = custom_load_hook
if custom_save_hook is not None:
self.custom_save_hooks[name] = custom_save_hook
def add_recoverables(self, recoverables):
"""Update the recoverables dict from the given mapping.
Arguments
---------
recoverables : mapping
Objects to recover.
They need a (unique) name: this is used to
connect the parameters in a checkpoint to the correct
recoverable. The name is also used in the filename of the
savefile for the objects parameters.
"""
if isinstance(recoverables, collections.abc.Mapping):
self.recoverables.update(recoverables)
else:
rec = repr(recoverables) # noqa: F841, rec is used in MSG
MSG = f"Checkpointer needs a mapping (e.g. dict), \
got {rec} instead."
raise AttributeError(MSG)
def save_checkpoint(
self, meta={}, end_of_epoch=True, name=None, verbosity=logging.INFO
):
"""Saves a checkpoint.
The whole checkpoint becomes a directory.
Saves each registered object's parameters in a separate file.
Also a meta file is added. The meta file by default has just the
unixtime (seconds since unix epoch), but you can add anything
relevant yourself. The meta information is later used to pick the
checkpoint to load.
The value of end_of_epoch is saved in the meta. This can affect how
epoch counters and dataset iterators load their state.
For multi-process saving there are cases where we may want to run
saving code on multiple processes (e.g. FSDP where we need to collect
parameters before saving). This works by creating a save folder
on the main process and communicating it to all processes, and then
letting each saver/loader method control whether it should save
on one or all processes.
Arguments
---------
meta : mapping, optional
A mapping which is added to the meta file in the checkpoint. The
key "unixtime" is included by default.
end_of_epoch : bool, optional
Whether the checkpoint is at the end of an epoch. True by default.
May affect loading.
name : str, optional
Specify a custom name for your checkpoint.
The name will still have a prefix added. If no name is given,
a name is created from a timestamp and a random unique id.
verbosity : logging level
Set logging level this save.
Returns
-------
Checkpoint
namedtuple [see above], the saved checkpoint, unless this is run
on a non-main process, in which case it returns None.
"""
ckpt_dir = None
if if_main_process():
if name is None:
ckpt_dir = self._new_checkpoint_dirpath()
else:
ckpt_dir = self._custom_checkpoint_dirpath(name)
os.makedirs(ckpt_dir, exist_ok=True)
saved_meta = self._save_checkpoint_metafile(
ckpt_dir / METAFNAME, meta, end_of_epoch
)
# Communicate ckpt_dir to all procs
ckpt_dir = ddp_broadcast(ckpt_dir, src=0)
saved_paramfiles = {}
for name, obj in self.recoverables.items():
objfname = f"{name}" + PARAMFILE_EXT
savepath = ckpt_dir / objfname
saved_paramfiles[name] = savepath
# First see if object has custom save hook:
if name in self.custom_save_hooks:
self.custom_save_hooks[name](obj, savepath)
continue
# Otherwise find the default saver for that type:
default_hook = get_default_hook(obj, DEFAULT_SAVE_HOOKS)
if default_hook is not None:
default_hook(obj, savepath)
continue
# If we got here, no custom hook or registered default hook
MSG = f"Don't know how to save {type(obj)}. Register default hook \
or add custom hook for this object."
raise RuntimeError(MSG)
if if_main_process():
ckpt_type = "end-of-epoch" if end_of_epoch else "intra-epoch"
logger.log(
verbosity, f"Saved an {ckpt_type} checkpoint in {ckpt_dir}"
)
return Checkpoint(ckpt_dir, saved_meta, saved_paramfiles)
# Explicitly return None if this is not the main process
return None
def save_and_keep_only(
self,
meta={},
end_of_epoch=True,
name=None,
num_to_keep=1,
keep_recent=True,
importance_keys=[],
max_keys=[],
min_keys=[],
ckpt_predicate=None,
verbosity=logging.INFO,
):
"""Saves a checkpoint, then deletes the least important checkpoints.
Essentially this combines ``save_checkpoint()`` and
``delete_checkpoints()`` in one call, providing short syntax.
Arguments
---------
meta : mapping, optional
A mapping which is added to the meta file in the checkpoint. The
key "unixtime" is included by default.
end_of_epoch : bool, optional
Whether the checkpoint is at the end of an epoch. True by default.
May affect loading.
name : str, optional
Specify a custom name for your checkpoint.
The name will still have a prefix added. If no name is given,
a name is created from a timestamp and a random unique id.
num_to_keep : int, optional
Number of checkpoints to keep. Defaults to 1. This deletes all
checkpoints remaining after filtering. Must be >=0.
keep_recent : bool, optional
Whether to keep the most recent ``num_to_keep`` checkpoints.
importance_keys : list, optional
A list of key functions used in sorting (see the sorted built-in).
Each callable defines a sort order and num_to_keep checkpoints are
kept for callable. The checkpoint with the highest keys are kept.
The functions are passed Checkpoint namedtuples (see above).
max_keys : list, optional
A list of keys for which the *highest* value will be kept.
min_keys : list, optional
A list of keys for which the *lowest* value will be kept.
ckpt_predicate : callable, optional
Use this to exclude some checkpoints from deletion. Before any
sorting, the list of checkpoints is filtered with this predicate.
Only the checkpoints for which ckpt_predicate is True can be
deleted. The function is called with Checkpoint namedtuples
(see above).
verbosity : int
The logging level, default logging.INFO
Note
----
Unlike save_checkpoint, this does not return anything, since we cannot
guarantee that the saved checkpoint actually survives deletion.
"""
self.save_checkpoint(
meta=meta, end_of_epoch=end_of_epoch, name=name, verbosity=verbosity
)
if keep_recent:
importance_keys.append(ckpt_recency)
self.delete_checkpoints(
num_to_keep=num_to_keep,
max_keys=max_keys,
min_keys=min_keys,
importance_keys=importance_keys,
ckpt_predicate=ckpt_predicate,
verbosity=verbosity,
)
def find_checkpoint(
self,
importance_key=None,
max_key=None,
min_key=None,
ckpt_predicate=None,
):
"""Picks a particular checkpoint from all available checkpoints.
If none of ``importance_key``, ``max_key``, and ``min_key`` is
used, then most recent checkpoint will be returned. No more than
one of them may be used.
Most functionality is actually implemented in ``find_checkpoints()``
but this is kept as a useful interface.
Arguments
---------
importance_key : callable, optional
The key function used in sorting.
The checkpoint with the highest returned value is picked.
The function is called with Checkpoint namedtuples.
max_key : str, optional
The checkpoint with the highest value for this key will
be returned. Only checkpoints with this key will be considered!
min_key : str, optional
The checkpoint with the lowest value for this key will
be returned. Only checkpoints with this key will be considered!
ckpt_predicate : callable, optional
Before sorting, the list of
checkpoints is filtered with this predicate.
See the filter builtin.
The function is called with Checkpoint namedtuples (see above).
By default, all checkpoints are considered.
Returns
-------
Checkpoint
If found.
None
If no Checkpoints exist/remain after filtering.
"""
ckpts_found = self.find_checkpoints(
importance_key=importance_key,
max_key=max_key,
min_key=min_key,
ckpt_predicate=ckpt_predicate,
max_num_checkpoints=None,
)
if ckpts_found:
return ckpts_found[0]
else:
return None
def find_checkpoints(
self,
importance_key=None,
max_key=None,
min_key=None,
ckpt_predicate=None,
max_num_checkpoints=None,
):
"""Picks multiple checkpoints.
If none of ``importance_key``, ``max_key``, and ``min_key`` is
used, then the most recent checkpoints will be returned. No more than
one of these may be used.
Arguments
---------
importance_key : callable, optional
The key function used in sorting.
The checkpoint with the highest returned value is picked.
The function is called with Checkpoint namedtuples.
max_key : str, optional
The checkpoint with the highest value for this key will
be returned. Only checkpoints with this key will be considered!
min_key : str, optional
The checkpoint with the lowest value for this key will
be returned. Only checkpoints with this key will be considered!
ckpt_predicate : callable, optional
Before sorting, the list of
checkpoints is filtered with this predicate.
See the filter builtin.
The function is called with Checkpoint namedtuples (see above).
By default, all checkpoints are considered.
max_num_checkpoints : int, None
The maximum number of checkpoints to return, or None to return all
found checkpoints.
Returns
-------
list
List containing at most the max specified number of Checkpoints.
"""
if importance_key is None and min_key is None and max_key is None:
importance_key = ckpt_recency
if max_key and not importance_key:
def importance_key(ckpt):
"""Defines the importance key."""
return ckpt.meta[max_key]
def ckpt_predicate(ckpt, old_predicate=ckpt_predicate):
"""Checkpoints predicate."""
if old_predicate is not None:
return max_key in ckpt.meta and old_predicate(ckpt)
else:
return max_key in ckpt.meta
elif min_key and not importance_key:
def importance_key(ckpt):
"""Defines the importance key."""
return -ckpt.meta[min_key]
def ckpt_predicate(ckpt, old_predicate=ckpt_predicate):
"""Checkpoints predicate."""
if old_predicate is not None:
return min_key in ckpt.meta and old_predicate(ckpt)
else:
return min_key in ckpt.meta
elif min_key or max_key:
raise ValueError(
"Must specify only one of 'importance_key', 'max_key', "
"and 'min_key'."
)
ckpts = self.list_checkpoints()
ckpts = list(filter(ckpt_predicate, ckpts))
# First sort by recency, so that importance being equal,
# the most checkpoints are returned
ckpts = sorted(ckpts, key=ckpt_recency, reverse=True)
if ckpts:
ranked_ckpts = sorted(ckpts, key=importance_key, reverse=True)
# NOTE: apparently, you can also slice [:None],
# and this is the same as [:], so the following if-else is not
# strictly speaking needed. However, this feature does not seem to
# be documented Python so I don't want to trust it.
if max_num_checkpoints is not None:
return ranked_ckpts[:max_num_checkpoints]
else: # No max number -> return all ckpts, but just sorted
return ranked_ckpts
else:
return [] # Be explicit :)
def recover_if_possible(
self,
importance_key=None,
max_key=None,
min_key=None,
ckpt_predicate=None,
):