Spaces:
Build error
Build error
import logging | |
import os | |
from urllib.parse import urlparse | |
try: | |
import comet_ml | |
except (ModuleNotFoundError, ImportError): | |
comet_ml = None | |
import yaml | |
logger = logging.getLogger(__name__) | |
COMET_PREFIX = 'comet://' | |
COMET_MODEL_NAME = os.getenv('COMET_MODEL_NAME', 'yolov5') | |
COMET_DEFAULT_CHECKPOINT_FILENAME = os.getenv('COMET_DEFAULT_CHECKPOINT_FILENAME', 'last.pt') | |
def download_model_checkpoint(opt, experiment): | |
model_dir = f'{opt.project}/{experiment.name}' | |
os.makedirs(model_dir, exist_ok=True) | |
model_name = COMET_MODEL_NAME | |
model_asset_list = experiment.get_model_asset_list(model_name) | |
if len(model_asset_list) == 0: | |
logger.error(f'COMET ERROR: No checkpoints found for model name : {model_name}') | |
return | |
model_asset_list = sorted( | |
model_asset_list, | |
key=lambda x: x['step'], | |
reverse=True, | |
) | |
logged_checkpoint_map = {asset['fileName']: asset['assetId'] for asset in model_asset_list} | |
resource_url = urlparse(opt.weights) | |
checkpoint_filename = resource_url.query | |
if checkpoint_filename: | |
asset_id = logged_checkpoint_map.get(checkpoint_filename) | |
else: | |
asset_id = logged_checkpoint_map.get(COMET_DEFAULT_CHECKPOINT_FILENAME) | |
checkpoint_filename = COMET_DEFAULT_CHECKPOINT_FILENAME | |
if asset_id is None: | |
logger.error(f'COMET ERROR: Checkpoint {checkpoint_filename} not found in the given Experiment') | |
return | |
try: | |
logger.info(f'COMET INFO: Downloading checkpoint {checkpoint_filename}') | |
asset_filename = checkpoint_filename | |
model_binary = experiment.get_asset(asset_id, return_type='binary', stream=False) | |
model_download_path = f'{model_dir}/{asset_filename}' | |
with open(model_download_path, 'wb') as f: | |
f.write(model_binary) | |
opt.weights = model_download_path | |
except Exception as e: | |
logger.warning('COMET WARNING: Unable to download checkpoint from Comet') | |
logger.exception(e) | |
def set_opt_parameters(opt, experiment): | |
"""Update the opts Namespace with parameters | |
from Comet's ExistingExperiment when resuming a run | |
Args: | |
opt (argparse.Namespace): Namespace of command line options | |
experiment (comet_ml.APIExperiment): Comet API Experiment object | |
""" | |
asset_list = experiment.get_asset_list() | |
resume_string = opt.resume | |
for asset in asset_list: | |
if asset['fileName'] == 'opt.yaml': | |
asset_id = asset['assetId'] | |
asset_binary = experiment.get_asset(asset_id, return_type='binary', stream=False) | |
opt_dict = yaml.safe_load(asset_binary) | |
for key, value in opt_dict.items(): | |
setattr(opt, key, value) | |
opt.resume = resume_string | |
# Save hyperparameters to YAML file | |
# Necessary to pass checks in training script | |
save_dir = f'{opt.project}/{experiment.name}' | |
os.makedirs(save_dir, exist_ok=True) | |
hyp_yaml_path = f'{save_dir}/hyp.yaml' | |
with open(hyp_yaml_path, 'w') as f: | |
yaml.dump(opt.hyp, f) | |
opt.hyp = hyp_yaml_path | |
def check_comet_weights(opt): | |
"""Downloads model weights from Comet and updates the | |
weights path to point to saved weights location | |
Args: | |
opt (argparse.Namespace): Command Line arguments passed | |
to YOLOv5 training script | |
Returns: | |
None/bool: Return True if weights are successfully downloaded | |
else return None | |
""" | |
if comet_ml is None: | |
return | |
if isinstance(opt.weights, str): | |
if opt.weights.startswith(COMET_PREFIX): | |
api = comet_ml.API() | |
resource = urlparse(opt.weights) | |
experiment_path = f'{resource.netloc}{resource.path}' | |
experiment = api.get(experiment_path) | |
download_model_checkpoint(opt, experiment) | |
return True | |
return None | |
def check_comet_resume(opt): | |
"""Restores run parameters to its original state based on the model checkpoint | |
and logged Experiment parameters. | |
Args: | |
opt (argparse.Namespace): Command Line arguments passed | |
to YOLOv5 training script | |
Returns: | |
None/bool: Return True if the run is restored successfully | |
else return None | |
""" | |
if comet_ml is None: | |
return | |
if isinstance(opt.resume, str): | |
if opt.resume.startswith(COMET_PREFIX): | |
api = comet_ml.API() | |
resource = urlparse(opt.resume) | |
experiment_path = f'{resource.netloc}{resource.path}' | |
experiment = api.get(experiment_path) | |
set_opt_parameters(opt, experiment) | |
download_model_checkpoint(opt, experiment) | |
return True | |
return None | |