MCPcopy Create free account
hub / github.com/tdrussell/diffusion-pipe / process_step

Method process_step

utils/saver.py:146–177  ·  view source on GitHub ↗
(self, step, examples)

Source from the content-addressed store, hash-verified

144 return epoch, checkpointed, saved
145
146 def process_step(self, step, examples):
147 checkpointed, saved = False, False
148 # Look at some simple "signal files" the user can write to save and optionally quit manually
149 should_manually_save = False
150 should_manually_quit = False
151 save_signal_file = self.save_root / 'save'
152 save_quit_signal_file = self.save_root / 'save_quit'
153 if save_signal_file.exists() and save_signal_file.is_file():
154 should_manually_save = True
155 dist.barrier()
156 if is_main_process():
157 os.remove(save_signal_file)
158 elif save_quit_signal_file.exists() and save_quit_signal_file.is_file():
159 should_manually_save = True
160 should_manually_quit = True
161 dist.barrier()
162 if is_main_process():
163 os.remove(save_quit_signal_file)
164
165 if 'save_every_n_steps' in self.config and step % self.config['save_every_n_steps'] == 0:
166 self.save_model(f'step{step}')
167 saved = True
168
169 if need_to_checkpoint(self.config) or should_manually_save:
170 self.save_checkpoint(step, examples)
171 checkpointed = True
172
173 if should_manually_quit:
174 print('Manually quitting')
175 sys.exit()
176
177 return checkpointed, saved

Callers 1

train.pyFile · 0.80

Calls 4

save_modelMethod · 0.95
save_checkpointMethod · 0.95
is_main_processFunction · 0.90
need_to_checkpointFunction · 0.85

Tested by

no test coverage detected