def test_set_dir(self):
temp_dir = tempfile.gettempdir()
hub.set_dir(temp_dir)
hub_model = hub.load('pytorch/vision', 'resnet18', pretrained=True)
self.assertEqual(self.resnet18_pretrained, hub_model.state_dict())
assert os.path.exists(temp_dir + '/vision_master')
shutil.rmtree(temp_dir + '/vision_master')
def test_set_dir(self):
temp_dir = tempfile.gettempdir()
hub.set_dir(temp_dir)
hub_model = hub.load('pytorch/vision', 'resnet18', pretrained=True)
self.assertEqual(sum_of_model_parameters(hub_model),
SUM_OF_PRETRAINED_RESNET18_PARAMS)
assert os.path.exists(temp_dir + '/pytorch_vision_master')
shutil.rmtree(temp_dir + '/pytorch_vision_master')
def test_set_dir(self):
temp_dir = tempfile.gettempdir()
hub.set_dir(temp_dir)
hub_model = hub.load('ailzhang/torchhub_example',
'mnist',
pretrained=True,
verbose=False)
self.assertEqual(sum_of_state_dict(hub_model.state_dict()),
SUM_OF_HUB_EXAMPLE)
assert os.path.exists(temp_dir + '/ailzhang_torchhub_example_master')
shutil.rmtree(temp_dir + '/ailzhang_torchhub_example_master')
def test_set_dir(self):
temp_dir = tempfile.gettempdir()
hub.set_dir(temp_dir)
hub_model = hub.load("pytorch/vision",
"resnet18",
weights="DEFAULT",
progress=False)
assert sum_of_model_parameters(hub_model).item() == pytest.approx(
SUM_OF_PRETRAINED_RESNET18_PARAMS)
assert os.path.exists(temp_dir + "/pytorch_vision_master")
shutil.rmtree(temp_dir + "/pytorch_vision_master")
def test_set_dir(self):
temp_dir = tempfile.gettempdir()
hub.set_dir(temp_dir)
hub_model = hub.load('pytorch/vision',
'resnet18',
pretrained=True,
progress=False)
assert sum_of_model_parameters(hub_model).item() == pytest.approx(
SUM_OF_PRETRAINED_RESNET18_PARAMS)
assert os.path.exists(temp_dir + '/pytorch_vision_master')
shutil.rmtree(temp_dir + '/pytorch_vision_master')
import warnings
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
from torch import hub, optim
from torch.backends import cudnn
from config import Config
from training import models
from training.datasets import get_dataloader
from training.tools.train_utils import parse_args, train, validate
CONFIG = Config()
hub.set_dir(CONFIG['TORCH_HOME'])
torch.backends.cudnn.benchmark = True
def main_worker(gpu, ngpus_per_node, args):
args.gpu = gpu
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
if args.distributed:
if args.dist_url == "env://" and args.rank == -1:
args.rank = int(os.environ["RANK"])
if args.multiprocessing_distributed:
# For multiprocessing distributed training, rank needs to be the
import argparse
from pathlib import Path
from torch import hub
from torchvision import models
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("model_arch", type=str, help="Architecture of the torchvision model")
parser.add_argument("torch_hub", type=str, help="Path to set as torch hub")
args = parser.parse_args()
hub_path = Path(args.torch_hub)
hub.set_dir(str(hub_path.resolve()))
model = getattr(models, args.model_arch)(pretrained=True)
from typing import List, Tuple
from PIL.Image import Image
from torch import hub
import whatplane.models.model_helpers as mh
from whatplane.models.predict_model import predict_image_data
BASE_DIR = Path(".")
MODELS_DIR = BASE_DIR / "models"
with open(BASE_DIR / "api/imagenet_class_index.json") as f:
imagenet_class_index = json.load(f)
# Set PyTorch cache directory to where the model is stored
hub.set_dir(str(MODELS_DIR.resolve()))
imagenet_model = mh.initialize_model(
"densenet161",
[item[1] for item in list(imagenet_class_index.values())],
replace_classifier=False,
)
whatplane_model = mh.load_model(MODELS_DIR / "model.pth")
def should_predict_whatplane(imagenet_probs: List[float],
imagenet_classes: List[str]) -> bool:
"""Function to check which model to use for prediction.
Args:
imagenet_probs (List[float]): List of probabilities returned from ImageNet
imagenet_classes (List[str]): List of class names returned from ImageNet
