Package: cresset
cressetcressetmakes it super simple to construct Pytorch modules directly from config files.
You no longer have to rely on the use of the unsafepicklemodule which can run arbitrary code.
Rather,cressetbuilds on the power ofconfactoryto enable loading the full model definition directly from JSON files.Simple ExampleHere's a simple example demonstrating how to create the same model from thePytorch Quickstart Tutorial.
First, write aquickstart.jsonfile with the following structure.{"type":"Sequential","args":[{"type":"Linear","in_features":784,"out_features":512},{"type":"ReLU"},{"type":"Linear","in_features":512,"out_features":512},{"type":"ReLU"},{"type":"Linear","in_features":512,"out_features":10}]}Now in Python you can build theModuledirectly from the config file:fromcresset.nnimportModulemodel=Module.from_config("quickstart.json")print(model)A More Complex ExampleLarge models have lots of layers and might even have conditional configuration options.
Using JSON directly would be combersome for defining such models.
Fortunately,cressetsupportsJsonnet, making configuration a breeze.
Below you can see a much more complicated example of creating a Transformer-like model from builtin Pytorch modules.
First create a file namedtransformer-like.jsonnetwith the following configuration:# An attention-less Transformer-like model (only performs feedforward ops)function(vocab_size,heads=8,dim=512,hidden_dim=2048,layers=6,dropout_p=0.1){# Make a reference to this transformer so nested objects can reference the top-level helperslocaltransformer=self,# Define a number of hidden helpers that can be overridden by callers of this function"dropout"::function(module,name="module"){"type":"Sequential","args":{[name]:module,"dropout":{"type":"Dropout","p":dropout_p}}},"residual"::function(name,module){"type":"Sequential","args":{"residual":{"type":"Parallel","args":[{"type":"Identity"},transformer.dropout(module,name=name)]},"sum":{"type":"Reduce","op":"SUM"},"norm":{"type":"LayerNorm","normalized_shape":dim}}},"ffn"::{"type":"Sequential","args":{"hidden":{"type":"Linear","in_features":dim,"out_features":hidden_dim},"activation":{"type":"ReLU","inplace":true},"output":{"type":"Linear","in_features":hidden_dim,"out_features":dim}}},"type":"Sequential","args":[{"type":"Embedding","embedding_dim":dim,"num_embeddings":vocab_size,}]+std.repeat([transformer.residual("ffn",transformer.ffn)],layers),}The top-level Jsonnet function has one required parameter,vocab_size, which must be specified.
This can be done using thebindingsparameter ofModule.from_config.
See the example Python below:fromcresset.nnimportModule# vocab_size is requiredModule.from_config("transformer-like.jsonnet",bindings={"vocab_size":65535})# can override default values, e.g. disable dropout for inferenceModule.from_config("transformer-like.jsonnet",bindings={"vocab_size":65535,"dropout_p":0})Additional FeaturesIn addition to support for Pytorch'storch.nn.Module,cressetalso wrapstorch.optim.Optimizerandtorch.optim.lr_scheduler.LRScheduler, allowing for defining training configurations.
Furthermore, it creates wrappers for the initialization functions intorch.nn.init, such that they can be used to initialize your model.
For example, the below code will run a set of initializers on your model parameters which match a regular expression.importloggingimportwarningsfromtypingimportList,Setimportattrfromconfactoryimportconfigurablefromcresset.nnimportSequentialfromcresset.nn.initimportInitializer,ZerosclassUninitializedParameters(Warning):"""A warning indicating an uninitialized parameter"""def__init__(self,parameters:List[str]):super().__init__(", ".join(parameters))@configurableclassModelWithInitializers(Sequential):"""A sequence model"""initializers:List[Initializer]=attr.ib(# The mypy attrs plugin has a bug, so we have to ignore the arg type# issue below, see https://github.com/python/mypy/issues/5313default=[Zeros(regex=r".*")],# type: ignore[arg-type]validator=attr.validators.deep_iterable(attr.validators.instance_of(Initializer)),)warn_uninitialized_params:bool=Falselog_initializers:bool=Falsedef__attrs_post_init__(self):matched:Set[str]=set()all_params:Set[str]=set()forname,parameterinself.named_parameters():all_params.add(name)forinitializerinself.initializers:ifinitializer.match(name):ifself.log_initializers:logging.info(f"Running initializer{initializer}for{name}")initializer(parameter)matched.add(name)uninitialized=all_params-matchedifself.warn_uninitialized_paramsanduninitialized:warnings.warn(UninitializedParameters(sorted(uninitialized)))Now if we use the previously defined Transformer-like model as a base, we can create an initialized version.
Name the new Jsonnet filetransformer-like-initialized.jsonnetwith the following contents:# Imports
local transformer = import "transformer-like.jsonnet";

function(vocab_size, log=false, warn_uninitialized=false)
  transformer(vocab_size, layers=1) {
    "type": "ModelWithInitializers",
    "initializers": [
      {"type": "Zeros", "regex": "bias"},
      {"type": "Ones", "regex": "norm.weight"},
      {"type": "XavierUniform", "regex": "ffn.hidden.weight", "gain": "linear"},
      {"type": "XavierUniform", "regex": "ffn.output.weight", "gain": "relu"}
    ],
    "log_initializers": log,
    "warn_uninitialized_params": warn_uninitialized,
  };Now we can load it similarly to the previous model, but this time the initializers will have been run:fromcresset.nnimportModulelogging.basicConfig(level=logging.INFO)Module.from_config("transformer-like-initialized.jsonnet",bindings={"vocab_size":65535,log:True})After running the above code, you should see output like this:INFO:root:Running initializer Ones(regex=re.compile('norm.weight')) for 0.norm.weight
INFO:root:Running initializer Zeros(regex=re.compile('bias')) for 0.norm.bias
INFO:root:Running initializer XavierUniform(regex=re.compile('ffn.hidden.weight'), gain=1) for 0.residual.1.ffn.hidden.weight
INFO:root:Running initializer Zeros(regex=re.compile('bias')) for 0.residual.1.ffn.hidden.bias
INFO:root:Running initializer XavierUniform(regex=re.compile('ffn.output.weight'), gain=1.4142135623730951) for 0.residual.1.ffn.output.weight
INFO:root:Running initializer Zeros(regex=re.compile('bias')) for 0.residual.1.ffn.output.biasLicenseThis repository is licensed under the MIT license.SPDX-License-Identifer: MIT