ModuleValidator.fix does not immediately change the model - how to deal with it?

username: general
cooked: <p>Hi,</p>
<p>I would like to use opacus with a DenseNet121. Since this neural network is by default not compatible with opacus, I need to use ModuleValidator.fix(model) before training the model.</p>
<p>Using the debugger, I noticed that ModuleValidator.fix(model) does not immediately change the state_dict of the model, but it is only changed later (probably after the training). Is this correct? Is there some way to change it immediately? (question 1)<br>
Specifically, it appears that the keys in the state dict later have an “_module.” substring before their original name.</p>
<p>Based on this, I have two more questions:</p>
<p>Q2: What is the best way to load a model that was trained using opacus? Normally, I would create a densenet object and then use model.load_state_dict(…). This, however, doesn’t work when the stored state dict is from a densenet121 trained with opacus, since the key names in the state dict are different. I can also not just use ModuleValidator.fix(model) before loading the state dict, since it does not immediately change the keys.</p>
<p>Q3: Could I just iterate over the keys key_name of the normal densenet121 model state dict, and then replace the values by the model_fixed_through_modelvalidator.state_dict()["_module."+key_name] or would this lead to “wrong” values in the loaded model (e.g., when BatchNorm is replaced with GroupNorm)</p>
<p>Thanks in advance!</p>

username: pstock
cooked: <p>Hey general, thanks for reaching out! Below are answers regarding your three questions:</p>
<ul>
<li>Q1: Yes, it’s indended to be used before training starts (or before checkpoint loading). Do you have a specific use case in mind that would support your suggestion (changing immediately the state dict)?</li>
<li>Q2: I would recommend that, during checkpointing when training with Opacus, you save <code>model._module.state_dict()</code>. This way, (1) you can load the checkpoint in a regular training loop as usual and (2) if you resume Opacus training from this checkpoint, you should call <code>model._module.load_state_dict()</code> after <code>make_private</code>.</li>
<li>Q3: See my geenric remark below.</li>
</ul>
<p>On your specific case, I am right to assume that you could switch batch norm layers by group norm layers <strong>before</strong> applying Opacus and <strong>then</strong> proceed as usual?</p>
<p>Hope this helps,<br>
Pierre</p>

username: general
cooked: <p>Hi Pierre, Thanks for your reply. To illustrate my problem, please have a look at the following code:</p>
<pre><code class="lang-auto">import torch
from torchvision import  models
from torch import nn
import numpy as np
import torchvision.transforms as transforms
from torchvision.datasets import CIFAR10
import torch.nn.functional as F
from opacus import PrivacyEngine
from opacus.validators import ModuleValidator

dataset = CIFAR10(root='data/', download=True, transform=transforms.ToTensor())
dataset = torch.utils.data.Subset(dataset,list(range(0,200)))
train_loader = torch.utils.data.DataLoader(dataset,batch_size=8, shuffle=True, num_workers=20, pin_memory=True)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
criterion = nn.BCELoss().to(device)

running_loss = 0.0

use_differential_privacy = True

if __name__ == "__main__":

    model = models.densenet121(pretrained=True)
    num_ftrs = model.classifier.in_features
    model.classifier = nn.Sequential(nn.Linear(num_ftrs, 10), nn.Sigmoid())

    if(use_differential_privacy):
        print("use differentialprivacy")
        model = ModuleValidator.fix(model)
        privacy_engine = PrivacyEngine()

    optimizer = torch.optim.SGD(model.parameters(),lr=0.01,momentum=0,weight_decay=0)
        
    if(use_differential_privacy):
        model, optimizer, train_loader = privacy_engine.make_private(
        module=model,
        optimizer=optimizer,
        data_loader=train_loader,
        noise_multiplier=1.1,
        max_grad_norm=1.0)

    model = model.to(device)
    
    print_freq = 1000
    running_loss = 0.0

    for i, data in enumerate(train_loader):
      imgs, labels = data

      batch_size = imgs.shape[0]
      imgs = imgs.to(device)
      labels = labels.to(device)

      
      labels = F.one_hot(labels, num_classes=10)
      labels = labels.to(torch.float32)

      optimizer.zero_grad()
      model.train()
      outputs = model(imgs)

      loss = criterion(outputs, labels)

      loss.backward()
      optimizer.step()  # update weights

      running_loss += loss * batch_size
      if (i % print_freq == 0):
        print(str(i * batch_size))

    epoch_loss_train = running_loss / len(dataset)

    if(use_differential_privacy):
        epsilon, best_alpha = privacy_engine.accountant.get_privacy_spent(
            delta=1/(10*len(dataset))
        )
        print("finished local training, privacy report:", privacy_engine.accountant.get_privacy_spent(delta=1/(10*len(train_loader))))

    print("finished training, now storing model")
    
    torch.save(model.state_dict(), "model.pt")
    
    model_load = models.densenet121(pretrained=True)
    num_ftrs = model_load.classifier.in_features
    model_load.classifier = nn.Sequential(nn.Linear(num_ftrs, 10), nn.Sigmoid())

    loaded_state = torch.load("model.pt")
    print("loading model")

    if(use_differential_privacy):
        model_load = ModuleValidator.fix(model_load)

    model_load.load_state_dict(loaded_state)
</code></pre>
<p>It results in the following error:</p>
<blockquote>
<p>Exception has occurred: RuntimeError<br>
Error(s) in loading state_dict for DenseNet:<br>
Missing key(s) in state_dict: “features.conv0.weight”, “features.norm0.weight”, “features.norm0.bias”, “features.denseblock1.denselayer1.norm1.weight”, “features.denseblock1.denselayer1.norm1.bias”, “features.denseblock1.denselayer1.conv1.weight”, “features.denseblock1.denselayer1.norm2.weight”, “features.denseblock1.denselayer1.norm2.bias”, “features.denseblock1.denselayer1.conv2.weight”, “features.denseblock1.denselayer2.norm1.weight”,<br>
…<br>
Unexpected key(s) in state_dict: “_module.features.conv0.weight”, “_module.features.norm0.weight”, “_module.features.norm0.bias”, “_module.features.denseblock1.denselayer1.norm1.weight”, “_module.features.denseblock1.denselayer1.norm1.bias”, “_module.features.denseblock1.denselayer1.conv1.weight”, “_module.features.denseblock1.denselayer1.norm2.weight”, “_module.features.denseblock1.denselayer1.norm2.bias”,<br>
…</p>
</blockquote>
<p>This happens, although I use the ModuleValidator before loading the new model. Do you know how I can fix my problem?</p>

username: pstock
cooked: <p>HI general,</p>
<p>Thanks for providing detailed context. You can simply load the state dict <em>after</em> you call <code>ModuleValidator.fix</code>. Could you please try this?</p>
<p>Thanks,<br>
Pierre</p>

username: general
cooked: <p>Hi Pierre,</p>
<p>Thanks for your response. I am not exactly sure what you mean, I call the “load_state_dict” function in the last line, after the ModuleValidator.fix.<br>
Or did you mean something different?<br>
Thanks,<br>
general</p>

username: pstock
cooked: <p>Hey general,</p>
<p>My bad, let me rephrase my suggestion.</p>
<ul>
<li>
<strong>Option 1</strong>. You can call <code>load_state_dict()</code> after calling <code>privacy_engine.make_private</code> (which you do not call here).</li>
<li>
<strong>Option 2</strong>. Replace <code>torch.save(model.state_dict(), "model.pt")</code> by <code>torch.save(model._module.state_dict(), "model.pt") </code> in case <code>use_differential_privacy</code> is <code>True</code>.</li>
</ul>
<p>Does that make sense?</p>
<p>PS: editing my comment to point out to a relevant <a href="https://github.com/pytorch/opacus/issues/435" rel="noopener nofollow ugc">planned enhancement</a>.</p>

username: general
cooked: <p>Hi Pierre,<br>
thanks, this helps - I especially like option 2.<br>
Best,<br>
general</p>