The PyTorch checkpoint example uses temporary files, which can be replaced with tempfile.
|
download_artifact( |
|
artifact_store=artifact_store, |
|
file_path=f"./tmp_model_{trial.number}.pt", |
|
artifact_id=artifact_id, |
|
) |
|
checkpoint = torch.load(f"./tmp_model_{trial.number}.pt") |
|
os.remove(f"./tmp_model_{trial.number}.pt") |
|
torch.save( |
|
{ |
|
"epoch": epoch, |
|
"model_state_dict": model.state_dict(), |
|
"optimizer_state_dict": optimizer.state_dict(), |
|
"accuracy": accuracy, |
|
}, |
|
f"./tmp_model_{trial.number}.pt", |
|
) |
|
artifact_id = upload_artifact( |
|
artifact_store=artifact_store, |
|
file_path=f"./tmp_model_{trial.number}.pt", |
|
study_or_trial=trial, |
|
) |
|
trial.set_user_attr("artifact_id", artifact_id) |
|
os.remove(f"./tmp_model_{trial.number}.pt") |