Package: jj-wingman
WingmanWingman is a Python library built for managing projects in deep learning.
Currently, Wingman is capable of automatically handling:Automatic weights saving and versioningSmooth WanbB integrationBuilding blocks for popular torch modulesSmoothly handling commandline arguments and*.yamlconfig filesInstallationpip3 install jj-wingmanWingman does not come with a PyTorch requirement by default, so all torch-functionality is by default disabled.
If you wish to use the PyTorch functionality of Wingman, install a compatible Pytorch version for your machine.
This is typically done via:pip3 install torchPhilosophyWingman is designed to be very horizontally integrated into your deep learning projects.
Instead of wrapping your models with wrappers or designing your models around Wingman, Wingman adapts to your projects, and you are the decider on when and where Wingman comes into play.ModulesWingman has several core modules and static functions, all of which can be used independently or in conjunction.from wingman import WingmanThis is the core module for Wingman, Wingman requires a bare minimumconfig.yamlfile somewhere accessible in your project directory, this can be generated by runningwingman-generate-yaml [optional filename]on the commandline.
A bare minimum yaml file is as follows:# wingman required paramsdebug:falseweights_directory:'weights'version_number:nullmark_number:0log_status:falseincrement:truelogging_interval:10max_skips:5greater_than:0.0wandb:falsewandb_name:''wandb_notes:''wandb_id:''wandb_entity:'jjshoots'wandb_project:'my_new_project'The parameters described are as follows:debug: Whether to set the model to debug modeweights_directory: Where should Wingman point to for model weight savingversion_number: Wingman versions different models using this number, if this is left as null, Wingman automatically chooses a number.mark_number: Wingman checkpoints models using this number if left at 0, automatic checkpoint numbering occurs if theincrementparameter is set to true.log_status: Whether to save a.txtlog file of Wingman's outputs.increment: Whether to increment mark number, if this is set to false, Wingman won't save multiple variations of the same model.logging_interval: During training, pass the training step to wingman, and afterlogging_intervalsteps has passed, Wingman will record the training. If Wingman has found a new lowest point, the model weights will be saved to a new file.max_skips: How many logging steps skipped without finding a new lowest point (specified using thelogging_intervalargument) before Wingman will save an intermediary checkpoint of the model.greater_than: You can tell Wingman to only checkpoint the model when the previous checkpointed loss is more than the current loss by this value.wandb: Whether to log things to WandB.wandb_name: The name of the run to be displayed in WandB. If left blank, Wingman automatically assigns one depending on the model version number.wandb_notes: Some helpful notes that you can leave for your runs that will be recorded in WandB.wandb_id: A Unique ID for this run. If left blank, one is automatically assigned.wandb_entity: Your username/organization name for WandB.wandb_project: The project that this run will be under.If you need more parameters, adding them to theconfig.yamlfile can be done freely, and these parameters can be easily referenced in the code later.
In addition, Wingman also automatically generates adeviceparameter, this parameter is automatically set to the GPU if your system has one, and can be used later for easy transfer of tensors.You can also use theconfig.yamlfile to define a set of sane defaults, and then override them using commandline arguments later.
This is because Wingman automatically converts all parameters in theconfig.yamlfile into commandline arguments using theargparsepython module.
This allows yourmain.pyfile to be much cleaner, without requiring the 50 odd lines of code at the top only for parameter definition.
For instance, if you havehack_level=10.3in theconfig.yaml, you can override its value via the commandline argument by doingpython3 your_program.py --hack_level=69.420.After defining yourconfig.yamlfile, the basic usage of Wingman is as follows:fromwingmanimportWingman# initialize Wingman and get all the parameters from the `config.yaml` filewm=Wingman("./config.yaml")cfg=wm.cfg# load the model and optimizer, `cfg.device` is automatically generatedmodel=Model(cfg.YOUR_PARAM,cfg.YOUR_OTHER_PARAM).to(cfg.device)optim=optimizer.AdamW(model.parameters(),lr=cfg.YOUR_LEARNING_RATE_PARAM,amsgrad=True)# we can check if we have trained this model before, and if we have, just load it# this checking is done using the `version_number` param, if `latest=True` is set,# Wingman automatically searches for the latest model checkpoint,# otherwise, Wingman uses the checkpoint specified by `mark_number`have_file,weight_file,optim_file=wm.get_weight_files(latest=True)ifhave_file:# Wingman simply returns a string of where the weight files are# no unnecessary wrapping!model.load_state_dict(torch.load(model_file))optim.load_state_dict(torch.load(optim_file))# let's run some training:while(training):...# training code hereloss=YOUR_LOSS_FUNCTION(model)...# when wandb is enabled (via `wandb = true` in the yaml file)# logging can be done by passing Wingman a dictionary# the dictionary is logged to wandb every logging intervalwm.log={"log":5,"these":2.34,"values":-5.432}# ... or you can force wandb logging usingwm.wandb_log()# let Wingman handle checkpointing for youupdate_weights,model_file,optim_file=wm.checkpoint(loss,training_step)ifupdate_weights:# if Wingman deems that the weights should be checkpointed, it returns# a string of where the weight files should go for it to be found latertorch.save(model.state_dict(),model_file)torch.save(optim.state_dict(),optim_file)Weights Directory CompressionIf you train a lot of models, the weights directory can start to use a lot of disk space.
Runningwingman-compress-weights [optional directory name]on the commandline will remove all redundant mark numbers and delete all empty files.from wingman import NeuralBlocksNeural blocks is a module for quickly prototyping neural network architectures.
It offers several easier methods of defining standardized modules:Simple 3-layer MLP with ReLU activation>>>fromwingmanimportNeuralBlocks>>>features=[4,16,64,3]>>>activation=["relu","tanh","identity"]>>>norm="batch">>>bias=True>>>MLP=NeuralBlocks.generate_linear_stack(features,activation,norm,bias)>>>print(MLP)Sequential((0):Sequential((0):BatchNorm1d(4,eps=1e-05,momentum=0.1,affine=True,track_running_stats=True)(1):Linear(in_features=4,out_features=16,bias=True)(2):ReLU())(1):Sequential((0):BatchNorm1d(16,eps=1e-05,momentum=0.1,affine=True,track_running_stats=True)(1):Linear(in_features=16,out_features=64,bias=True)(2):Tanh())(2):Sequential((0):BatchNorm1d(64,eps=1e-05,momentum=0.1,affine=True,track_running_stats=True)(1):Linear(in_features=64,out_features=3,bias=True)(2):Identity()))4 layer convolutional network with same padding and MaxPool2d across the last two layers>>>fromwingmanimportNeuralBlocks>>>channels=[3,16,32,64,1]>>>kernels=[3,1,3,1]>>>pooling=[0,0,2,2]>>>activation=["relu","relu","tanh","tanh"]>>>padding=None# this is equivalent to same padding>>>norm="batch">>>CNN=NeuralBlocks.generate_conv_stack(channels,kernels,pooling,activation,padding,norm)>>>print(CNN)Sequential((0):Sequential((0):BatchNorm2d(3,eps=1e-05,momentum=0.1,affine=True,track_running_stats=True)(1):Conv2d(3,16,kernel_size=(3,3),stride=(1,1),padding=(1,1))(2):ReLU())(1):Sequential((0):BatchNorm2d(16,eps=1e-05,momentum=0.1,affine=True,track_running_stats=True)(1):Conv2d(16,32,kernel_size=(1,1),stride=(1,1))(2):ReLU())(2):Sequential((0):BatchNorm2d(32,eps=1e-05,momentum=0.1,affine=True,track_running_stats=True)(1):Conv2d(32,64,kernel_size=(3,3),stride=(1,1),padding=(1,1))(2):MaxPool2d(kernel_size=2,stride=2,padding=0,dilation=1,ceil_mode=False)(3):Tanh())(3):Sequential((0):BatchNorm2d(64,eps=1e-05,momentum=0.1,affine=True,track_running_stats=True)(1):Conv2d(64,1,kernel_size=(1,1),stride=(1,1))(2):MaxPool2d(kernel_size=2,stride=2,padding=0,dilation=1,ceil_mode=False)(3):Tanh()))2 layer transposed convolutional network>>>fromwingmanimportNeuralBlocks>>>channels=[64,32,3]>>>kernels=[4,4]>>>padding=[1,1]>>>stride=[2,2]>>>activation=["lrelu","lrelu"]>>>norm="non">>>TCNN=NeuralBlocks.generate_deconv_stack(channels,kernels,padding,stride,activation,norm)>>>print(TCNN)Sequential((0):Sequential((0):ConvTranspose2d(64,32,kernel_size=(4,4),stride=(2,2),padding=(1,1))(1):LeakyReLU(negative_slope=0.01))(1):Sequential((0):ConvTranspose2d(32,3,kernel_size=(4,4),stride=(2,2),padding=(1,1))(1):LeakyReLU(negative_slope=0.01)))The Neural Blocks module also has functions that can generate single modules, refer to the file itself for more details.from wingman import ReplayBufferThis is a replay buffer designed around Torch's Dataloader class for reinforcement learning projects.
This allows easy bootstrapping of the Dataloader's excellent shuffling and pre-batching capabilities.
In addition, all the data is stored as a numpy array in a contiguous block of memory, allowing very fast retrieval.
ReplayBuffer also doesn't put any limits on tuple length per transition; some people prefer to store ${S, A, R, S'}$, some prefer to store ${S, A, R, S', A'}$ - ReplayBuffer doesn't care!
The length of the tuple can be as long or as short as you want, as long as every tuple fed in is the same length and each element of the tuple is the same shape.
There is no need to predefine the shape of the inputs that you want to put in the ReplayBuffer, it automatically infers the shape and computes memory usage upon the first tuple stored.
The basic usage of the ReplayBuffer is as follows:importtorchfromwingmanimportReplayBuffer# we define the replay buffer to be able to store 1000 tuples of informationmemory=ReplayBuffer(mem_size=1000)# get the first observation from the environmentnext_obs=env.reset()# iterate until the environment is completewhileenv.doneisFalse:# rollover the observationobs=next_obs# get an action from the policyact=policy(obs)# sample a new transitionnext_obs,rew,dne,next_lbl=env.step(act)# store stuff in the replay buffermemory.push((obs,act,rew,next_obs,dne))# perform training using the bufferdataloader=torch.utils.data.DataLoader(memory,batch_size=32,shuffle=True,drop_last=False)# easily treat the replay buffer as an iterable that we can iterate throughforbatch_num,stuffinenumerate(dataloader):observations=gpuize(stuff[0],"cuda:0")actions=gpuize(stuff[1],"cuda:0")rewards=gpuize(stuff[2],"cuda:0")next_states=gpuize(stuff[3],"cuda:0")dones=gpuize(stuff[4],"cuda:0")from wingman import gpuize, cpuizeThese are quality of life standalone functions.>>>importnumpyasnp>>>fromwingmanimportgpuize,cpuize>>>foo=np.random.random((5,5))>>>print(foo)[[0.283927640.509369830.554336160.456145180.82523046][0.774370720.209003820.862204940.690712390.94863786][0.700828650.927800180.983929650.769451650.72886401][0.477024850.549685220.221109420.724362760.42574472][0.783302210.848888370.685291670.618789020.13556213]]>>>bar=gpuize(foo,"cuda:0")>>>print(bar)tensor([[0.2839,0.5094,0.5543,0.4561,0.8252],[0.7744,0.2090,0.8622,0.6907,0.9486],[0.7008,0.9278,0.9839,0.7695,0.7289],[0.4770,0.5497,0.2211,0.7244,0.4257],[0.7833,0.8489,0.6853,0.6188,0.1356]],device='cuda:0')>>>baz=cpuize(bar)>>>print(baz)[[0.283927640.509369830.554336160.456145180.82523046][0.774370720.209003820.862204940.690712390.94863786][0.700828650.927800180.983929650.769451650.72886401][0.477024850.549685220.221109420.724362760.42574472][0.783302210.848888370.685291670.618789020.13556213]]If used in conjunction with Wingman, it's possible to simply dogpuize(foo, cfg.device)whencfgis defined as done earlier.Why does every plane have two pilots?
Really, you only need one pilot.
Let's take out the second pilot.
Let's let the bloody computer fly it.