###################################################################### # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # # SPDX-License-Identifier: MIT-0 # ###################################################################### import os import torch import torch.neuron from PIL import Image import torchvision from torchvision import transforms from common_settings import default_image_size, default_batch_size, default_model_name image_size = int(os.getenv('IMAGE_SIZE', default_image_size)) batch_size = int(os.getenv('BATCH_SIZE', default_batch_size)) # Setting up NeuronCore groups # This value should be 4 on inf1.xlarge and inf1.2xlarge, 16 for inf1.6xlarge num_cores = 4 nc_env = ','.join(['1'] * num_cores) os.environ['NEURONCORE_GROUP_SIZES'] = nc_env os.environ['TOKENIZERS_PARALLELISM'] = 'False' neuron_model_file = '%s_inf_%d_%d.pt'%(default_model_name,image_size, batch_size) img_cat = Image.open("data/cat.png").convert('RGB') # # Create a preprocessing pipeline # preprocess = transforms.Compose([ transforms.Resize(image_size), transforms.CenterCrop(image_size), transforms.ToTensor(), transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] )]) # # Pass the image for preprocessing and the image preprocessed # img_cat_preprocessed = preprocess(img_cat) img_cat_preprocessed_unsqueeze = torch.unsqueeze(img_cat_preprocessed, 0) batch_img_cat_tensor = torch.cat([img_cat_preprocessed_unsqueeze] * batch_size) model_ft = torchvision.models.resnet50(pretrained=True) model_ft.eval() # REmove None Attributes # https://forums.developer.nvidia.com/t/torchscripted-pytorch-lightning-module-fails-to-load/181002 remove_attributes = [] for key, value in vars(model_ft).items(): if value is None: remove_attributes.append(key) for key in remove_attributes: delattr(model_ft, key) orig_output = model_ft(batch_img_cat_tensor) print('Original Model Output:', orig_output) # Compile the model neuron_model = torch.neuron.trace(model_ft, batch_img_cat_tensor) #, # compiler_args = ['--neuroncore-pipeline-cores', str(num_cores)]) # Save the compiled model for later use neuron_model.save(neuron_model_file) # Load the saved model and perform inference neuron_model_reloaded = torch.jit.load(neuron_model_file) neuron_output = neuron_model_reloaded(batch_img_cat_tensor) print('Neuron Model Output:', neuron_output)