[FR][hub] load_state_dict_from_url should avoid using tempfile or allow specifying tmp dir

[ssnl]

🚀 Feature

The torch.hub.load_state_dict_from_url function

pytorch/torch/hub.py

Line 409 in 456e66d

def load_state_dict_from_url(url, model_dir=None, map_location=None, progress=True):

calls _download_url_to_file which creates a temp file and download to that first, does sha check, and them move to the target location.

pytorch/torch/hub.py

Line 381 in 456e66d

f = tempfile.NamedTemporaryFile(delete=False)

When the model is huge and /tmp is small, users can not actually finish downloading. And there is no way to specify the tmp location.

Alternatives

I see two potential solutions here:

1. Allow specifying tmp_folder
2. By default download to file_name + random_uuid, and use that as a temp file. This way we are always downloading to the same disk and will not have such issues.

[ssnl]

cc @ailzhang for comments!

[ailzhang]

Hi @ssnl I think I have a fix for it.

We can specify tempfile.TemporaryFile(delete=False, dir=dst) in _download_url_to_file and place the file directly in the dst folder. It seems much simpler this way :D

If this sounds good to you, I can quickly make a PR later today, although this might not go in 1.2 release :P

[ssnl]

@ailzhang That sounds a pretty elegant fix. Thanks so much for the quick reply! :)
There is no rush. I live on the edge and run my code on master builds anyways :D