151 lines
4.6 KiB
Python
151 lines
4.6 KiB
Python
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
|