} } def gen_classifier_loader(name, d): def classifier_loader(): if name == 'googlenet/inceptionv1': model = torch_models.__dict__[d['arch']](pretrained=False, aux_logits=False, transform_input=True) else: model = torch_models.__dict__[d['arch']](pretrained=False) load_model_state_dict(model, name) return model return classifier_loader for name, d in model_params.items(): registry.add_model( Model(name=name, arch=d['arch'], transform=StandardTransform(d['img_resize_size'], d['img_crop_size']), normalization=StandardNormalization(d['mean'], d['std'], d['input_space']), classifier_loader=gen_classifier_loader(name, d), eval_batch_size=d['eval_batch_size'], adversarial_batch_size=d['adversarial_batch_size'] if 'adversarial_batch_size' in d else None))
import torchvision.models as torch_models from registry import registry from models.model_base import Model, StandardTransform, StandardNormalization from mldb.utils import load_model_state_dict def classifier_loader(): model = torch_models.resnet50() load_model_state_dict(model, 'resnet50_augmix') return model registry.add_model( Model( name='resnet50_augmix', arch='resnet50', transform=StandardTransform(img_resize_size=256, img_crop_size=224), normalization=StandardNormalization(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), classifier_loader=classifier_loader, eval_batch_size=256, ))
'adversarial_batch_size': 1}, 'BiT-M-R101x3-ILSVRC2012': { 'arch': 'BiT-M-R101x3', 'eval_batch_size': 32, 'adversarial_batch_size': 1}, 'BiT-M-R152x4-ILSVRC2012': { 'arch': 'BiT-M-R152x4', 'eval_batch_size': 8}, } def gen_classifier_loader(name, d): def classifier_loader(): model = KNOWN_MODELS[d['arch']](head_size=1000) load_model_state_dict(model, name) return model return classifier_loader for name, d in model_params.items(): registry.add_model( Model( name = name, arch = d['arch'], transform = transforms.Compose([transforms.Resize((480, 480)), transforms.ToTensor()]), normalization = StandardNormalization([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]), classifier_loader = gen_classifier_loader(name, d), eval_batch_size = d['eval_batch_size'], adversarial_batch_size = d['adversarial_batch_size'] if 'adversarial_batch_size' in d else None, ) )
def classifier_loader(): model = EfficientNet.from_name(d['arch']) load_model_state_dict(model, name) return model return classifier_loader for name, d in model_params.items(): registry.add_model( Model(name=name, arch=d['arch'], transform=transforms.Compose([ transforms.Resize(d['img_size'] + CROP_PADDING, interpolation=PIL.Image.BICUBIC), transforms.CenterCrop(d['img_size']), transforms.ToTensor() ]), normalization=StandardNormalization(d['mean'], d['std']), classifier_loader=gen_classifier_loader(name, d), eval_batch_size=d['eval_batch_size'], adversarial_batch_size=d['adversarial_batch_size'] if 'adversarial_batch_size' in d else None)) def noisystudent_loader(): model = timm.create_model('tf_efficientnet_l2_ns', pretrained=False) load_model_state_dict(model, 'efficientnet-l2-noisystudent') return model registry.add_model(
def gen_classifier_loader(name, d): def classifier_loader(): model = torch_models.__dict__[d['arch']]() load_model_state_dict(model, name) model = Smooth(model, d['noise_sigma'], d['n'], d['alpha'], d['mean'], d['std']) return model return classifier_loader def classify(images, model, class_sublist, adversarial_attack): if adversarial_attack: images = pgd_style_attack(adversarial_attack, images, model) return model.predict_batch(images, class_sublist=class_sublist) for name, d in model_params.items(): registry.add_model( Model( name=name, arch=d['arch'], transform=StandardTransform(d['img_resize_size'], d['img_crop_size']), classifier_loader=gen_classifier_loader(name, d), eval_batch_size=d['eval_batch_size'], adversarial_batch_size=d['adversarial_batch_size'] if 'adversarial_batch_size' in d else None, classify=classify, ))