make strict and fallback_to_scratch args configurable to add support for partial checkpoint restore

ashors1 · ashors1 · commit 0678aa3e8b7d · 2023-09-13T08:52:31.000-07:00
diff --git a/t5x/train.py b/t5x/train.py
@@ -354,6 +354,11 @@ def _verify_matching_vocabs(cfg: utils.DatasetConfig):
               checkpoint_cfg.save and checkpoint_cfg.save.save_dataset
           ),
           state_transformation_fns=state_transforms_for_restore,
+          strict=(checkpoint_cfg.restore.strict
+                  if checkpoint_cfg.restore is not None else True
+          ),
+          fallback_to_scratch=(checkpoint_cfg.restore.fallback_to_scratch
+                               if checkpoint_cfg.restore is not None else False)
       )
   ]
   # 2. From a checkpoint specified by `checkpoint_cfg.restore.path`, if set.

Original file line number	Diff line number	Diff line change
`@@ -354,6 +354,11 @@ def _verify_matching_vocabs(cfg: utils.DatasetConfig):`
`354`	`354`	`checkpoint_cfg.save and checkpoint_cfg.save.save_dataset`
`355`	`355`	`),`
`356`	`356`	`state_transformation_fns=state_transforms_for_restore,`
	`357`	`+ strict=(checkpoint_cfg.restore.strict`
	`358`	`+ if checkpoint_cfg.restore is not None else True`
	`359`	`+ ),`
	`360`	`+ fallback_to_scratch=(checkpoint_cfg.restore.fallback_to_scratch`
	`361`	`+ if checkpoint_cfg.restore is not None else False)`
`357`	`362`	`)`
`358`	`363`	`]`
`359`	`364`	# 2. From a checkpoint specified by `checkpoint_cfg.restore.path`, if set.