RuntimeError: Expanded Weights encountered but cannot handle function mul

[heilrahc]

I'm trying to use privacy engine to achieve dp-private training. However, the model that will be used contained self customized layers which Opacus doesn't support like the following:
`class Conv2d(nn.Conv2d):
def init(self, in_channels, out_channels, kernel_size, stride=1,
padding=0, dilation=1, groups=1,
bias=True, padding_mode='zeros', name=None):
super(Conv2d, self).init(
in_channels, out_channels, kernel_size, stride, padding,
dilation, groups, bias, padding_mode)
self.name = name
self.register_buffer('weight_mask', torch.ones(self.weight.shape))
if self.bias is not None:
self.register_buffer('bias_mask', torch.ones(self.bias.shape))

def _conv_forward(self, input, weight, bias):
    if self.padding_mode != 'zeros':
        return F.conv2d(F.pad(input, self._padding_repeated_twice, mode=self.padding_mode),
                        weight, bias, self.stride,
                        _pair(0), self.dilation, self.groups)
    return F.conv2d(input, weight, bias, self.stride,
                    self.padding, self.dilation, self.groups)

def forward(self, input):
    W = self.weight * self.weight_mask
    if self.bias is not None:
        b = self.bias * self.bias_mask
    else:
        b = self.bias
    return self._conv_forward(input, W, b)`

I tried to use opacus with ExpandedWeights like this:

model = GradSampleModuleExpandedWeights(model) model, optimizer, train_loader = privacy_engine.make_private_with_epsilon( module=model, optimizer=optimizer, data_loader=train_loader, target_epsilon=sigma, target_delta=delta, epochs=args.post_epochs, max_grad_norm=clip, grad_sample_mode="ew")

but I got error:

File "/Layers/layers.py", line 52, in forward
W = self.weight * self.weight_mask
File "/home/mine01/Desktop/code/DP-SGD_SNIP/dpsgd_py/lib/python3.8/site-packages/torch/nn/utils/_expanded_weights/expanded_weights_impl.py", line 122, in torch_function
raise RuntimeError(f"Expanded Weights encountered but cannot handle function {func.name}")
RuntimeError: Expanded Weights encountered but cannot handle function mul

what's the proper way to use opacus privacy engine on my model?

[alexandresablayrolles]

Thanks for raising this issue, did you try other values for grad_sample_mode?

[HuanyuZhang]

Closing the issue due to long-time no reply. Please re-open if needed.