HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

HUB checkpoint upload fix (#10088)

Browse source 
Co-authored-by: Hassaan Farooq <103611273+hassaanfarooq01@users.noreply.github.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
Glenn Jocher 2024-04-16 14:52:03 -07:00 • committed by GitHub
parent 28f77d7011
commit eadb9cadf8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 18 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

19
ultralytics/hub/session.py
View file

@ -48,7 +48,7 @@ class HUBTrainingSession:
self.rate_limits = { self.rate_limits = {
"metrics": 3.0, "metrics": 3.0,
"ckpt": 900.0, "ckpt": 90.0,
"heartbeat": 300.0, "heartbeat": 300.0,
} # rate limits (seconds) } # rate limits (seconds)
self.metrics_queue = {} # holds metrics for each epoch until upload self.metrics_queue = {} # holds metrics for each epoch until upload
@ -213,6 +213,7 @@ class HUBTrainingSession:
thread=True, thread=True,
verbose=True, verbose=True,
progress_total=None, progress_total=None,
stream_reponse=None,
*args, *args,
**kwargs, **kwargs,
): ):
@ -232,6 +233,8 @@ class HUBTrainingSession:
if progress_total: if progress_total:
self._show_upload_progress(progress_total, response) self._show_upload_progress(progress_total, response)
elif stream_reponse:
self._iterate_content(response)
if HTTPStatus.OK <= response.status_code < HTTPStatus.MULTIPLE_CHOICES: if HTTPStatus.OK <= response.status_code < HTTPStatus.MULTIPLE_CHOICES:
# if request related to metrics upload # if request related to metrics upload
@ -335,6 +338,7 @@ class HUBTrainingSession:
timeout=3600, timeout=3600,
thread=not final, thread=not final,
progress_total=progress_total, progress_total=progress_total,
stream_reponse=True,
) )
else: else:
LOGGER.warning(f"{PREFIX}WARNING ⚠️ Model upload issue. Missing model {weights}.") LOGGER.warning(f"{PREFIX}WARNING ⚠️ Model upload issue. Missing model {weights}.")
@ -353,3 +357,16 @@ class HUBTrainingSession:
with TQDM(total=content_length, unit="B", unit_scale=True, unit_divisor=1024) as pbar: with TQDM(total=content_length, unit="B", unit_scale=True, unit_divisor=1024) as pbar:
for data in response.iter_content(chunk_size=1024): for data in response.iter_content(chunk_size=1024):
pbar.update(len(data)) pbar.update(len(data))
def _iterate_content(self, response: requests.Response) -> None:
"""
Process the streamed HTTP response data.
Args:
response (requests.Response): The response object from the file download request.
Returns:
None
"""
for data in response.iter_content(chunk_size=1024):
pass # Do nothing with data chunks