Package: modelscan
ModelScan: Protection Against Model Serialization AttacksMachine Learning (ML) models are shared publicly over the internet, within teams and across teams. The rise of Foundation Models have resulted in public ML models being increasingly consumed for further training/fine tuning. ML Models are increasingly used to make critical decisions and power mission-critical applications.
Despite this, models are not scanned with the rigor of a PDF file in your inbox.This needs to change, and proper tooling is the first step.ModelScan is an open source project that scans models to determine if they contain
unsafe code. It is the first model scanning tool to support multiple model formats.
ModelScan currently supports: H5, Pickle, and SavedModel formats. This protects you
when using PyTorch, TensorFlow, Keras, Sklearn, XGBoost, with more on the way.TL;DRIf you are ready to get started scanning your models, it is simple:pipinstallmodelscanWith it installed, scan a model:modelscan-p/path/to/model_file.h5Why You Should Scan ModelsModels are often created from automated pipelines, others may come from a data scientist’s laptop. In either case the model needs to move from one machine to another before it is used. That process of saving a model to disk is called serialization.AModel Serialization Attackis where malicious code is added to the contents of a model during serialization(saving) before distribution — a modern version of the Trojan Horse.The attack functions by exploiting the saving and loading process of models. When you load a model withmodel = torch.load(PATH), PyTorch opens the contents of the file and begins to running the code within. The second you load the model the exploit has executed.AModel Serialization Attackcan be used to execute:Credential Theft(Cloud credentials for writing and reading data to other systems in your environment)Data Theft(the request sent to the model)Data Poisoning(the data sent after the model has performed its task)Model Poisoning(altering the results of the model itself)These attacks are incredibly simple to execute and you can view working examples in our 📓notebooksfolder.Getting StartedHow ModelScan WorksIf loading a model with your machine learning framework automatically executes the attack,
how does ModelScan check the content without loading the malicious code?Simple, it reads the content of the file one byte at a time just like a string, looking for
code signatures that are unsafe. This makes it incredibly fast, scanning models in the time it
takes for your computer to process the total filesize from disk(seconds in most cases). It also secure.ModelScan ranks the unsafe code as:CRITICALHIGHMEDIUMLOWIf an issue is detected, reach out to the author's of the model immediately to determine the cause.In some cases, code may be embedded in the model to make things easier to reproduce as a data scientist, but
it opens you up for attack. Use your discretion to determine if that is appropriate for your workloads.What Models and Frameworks Are Supported?This will be expanding continually, so look out for changes in our release notes.At present, ModelScan supports any Pickle derived format and many others:ML LibraryAPISerialization Formatmodelscan supportPytorchtorch.save() and torch.load()PickleYesTensorflowtf.saved_model.save()Protocol BufferYesKeraskeras.models.save(save_format= 'h5')HD5 (Hierarchical Data Format)Yeskeras.models.save(save_format= 'keras')Keras V3 (Hierarchical Data Format)YesClassic ML Libraries (Sklearn, XGBoost etc.)pickle.dump(), dill.dump(), joblib.dump(), cloudpickle.dump()Pickle, Cloudpickle, Dill, JoblibYesInstallationModelScan is installed on your systems as a Python package(Python 3.8 to 3.11 supported). As shown from above you can install
it by running this in your terminal:pipinstallmodelscanTo include it in your project's dependencies so it is available for everyone, add it to yourrequirements.txtorpyproject.tomllike this:modelscan=">=0.1.1"Using ModelScan via CLIModelScan supports the following arguments via the CLI:UsageArgumentExplanationmodelscan -h-h or --helpView usage helpmodelscan -v-v or --versionView version informationmodelscan -p /path/to/model_file-p or --pathScan a locally stored modelmodelscan -p /path/to/model_file --settings-file ./modelscan-settings.toml--settings-fileScan a locally stored model using custom configurationsmodelscan create-settings-file-l or --locationCreate a configurable settings filemodelscan -r-r or --reporting-formatFormat of the output. Options are console,       json, or custom (to be defined in settings-file). Default is consolemodelscan -r reporting-format -o file-name-o or --output-fileOptional file name for output reportmodelscan --show-skipped--show-skippedPrint a list of files that were skipped      during the scanRemember models are just like any other form of digital media, you should scan content from any untrusted source before use.CLI Exit CodesThe CLI exit status codes are:0: Scan completed successfully, no vulnerabilities found1: Scan completed successfully, vulnerabilities found2: Scan failed, modelscan threw an error while scanning3: No supported files were passed to the tool4: Usage error, CLI was passed invalid or incomplete optionsUnderstanding The ResultsOnce a scan has been completed you'll see output like this if an issue is found:Here we have a model that has an unsafe operator for bothReadFileandWriteFilein the model.
Clearly we do not want our models reading and writing files arbitrarily. We would now reach out
to the creator of this model to determine what they expected this to do. In this particular case
it allows an attacker to read our AWS credentials and write them to another place.That is a firm NO for usage.Integrating ModelScan In Your ML Pipelines and CI/CD PipelinesAd-hoc scanning is a great first step, please drill it into yourself, peers, and friends to do
this whenever they pull down a new model to explore. It is not sufficient to improve security
for production MLOps processes.Model scanning needs to be performed more than once to accomplish the following:Scan all pre-trained models before loading it for further work to prevent a compromised
model from impacting your model building or data science environments.Scan all models after training to detect a supply chain attack that compromises new models.Scan all models before deploying to an endpoint to ensure that the model has not been compromised after storage.The red blocks below highlight this in a traditional ML Pipeline.The processes would be the same for fine-tuning or any modifications of LLMs, foundational models, or external model.Embed scans into deployment processes in your CI/CD systems to secure usage
as models are deployed as well if this is done outside your ML Pipelines.Diving DeeperInside the 📓notebooksfolder you can explore a number of notebooks that showcase
exactly how Model Serialization Attacks can be performed against various ML Frameworks like TensorFlow and PyTorch.To dig more into the meat of how exactly these attacks work check out 🖹Model Serialization Attack Explainer.If you encounter any other approaches for evaluating models in a static context, please reach out, we'd love
to learn more!LicensingCopyright 2023 Protect AILicensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License athttp://www.apache.org/licenses/LICENSE-2.0Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.AcknowledgementsWe were heavily inspired byMatthieu Maitrewho builtPickleScan.
We appreciate the work and have extended it significantly with ModelScan. ModelScan is OSS’ed in the similar spirit as PickleScan.ContributingWe would love to have you contribute to our open source ModelScan project.
If you would like to contribute, please follow the details onContribution page.