CNNs.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import resnet34, ResNet34_Weights
from torchvision.models import resnet18, ResNet18_Weights

from typing import Union, Any, Tuple


def load_model(model_name: str, 
               dataset_name: str, 
               device: Union[torch.device, str] = torch.device('cuda'),
               ) -> nn.Module:
    
    if 'cifar' in dataset_name:
        full_name = dataset_name + "_" + model_name
        model = torch.hub.load("chenyaofo/pytorch-cifar-models", full_name, pretrained=True)

    elif 'imagenet' in dataset_name:
        if model_name == "resnet34":
            model = resnet34(weights=ResNet34_Weights.DEFAULT).to(device)
        
        elif model_name == "resnet18":
            model = resnet18(weights=ResNet18_Weights.DEFAULT).to(device)

    return model.to(device)