Skip to content

Instantly share code, notes, and snippets.

@dhuynh95

dhuynh95/convert_bnn.py

Created November 28, 2019 18:55
Show Gist options
  • Save dhuynh95/db862a51fab8ee9c02853a4190d25e68 to your computer and use it in GitHub Desktop.
Save dhuynh95/db862a51fab8ee9c02853a4190d25e68 to your computer and use it in GitHub Desktop.
Download ZIP
Convert a NN to a BNN
Raw
convert_bnn.py
# Convert nn.Dropout to CustomDropout module
get_args = lambda dp : {"p" : dp.p}
convert_layers(learn.model,nn.Dropout,CustomDropout,get_args)
# Turn on the stochasticity, I use verbose just to make sure it's working fine
switch_custom_dropout(learn.model,True,verbose=True)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment