fix: check if data file exists before downloading
@@ -61,8 +61,11 @@ class UNetsTrial(TFKerasTrial):
# Use a file lock so only one worker on each node does the download
with filelock.FileLock(os.path.join(weights_dir, "download.lock")):
- urllib.request.urlretrieve(mobilenet_link, weights_dir + data_file)
- return weights_dir + data_file
+ full_weights_path = weights_dir + data_file
+ if not os.path.exists(full_weights_path):
+ urllib.request.urlretrieve(mobilenet_link, full_weights_path + ".part")
+ os.rename(full_weights_path + ".part", full_weights_path)
+ return full_weights_path
def build_model(self):
model_weights_loc = self.download_weights()
