Hi, I am enjoying using the opacus package to apply differential privacy to the training process of my models, I am struggling to get it to work with my TVAE implementation though, could someone let me know why I get an Incompatible Module Exception, I am using similar modules to in all my other generative models. See my code below:
import numpy as np
import torch
from torch.nn import Linear, Module, Parameter, ReLU, Sequential
from torch.nn.functional import cross_entropy
from torch.optim import Adam
from torch.utils.data import DataLoader, TensorDataset
from models.CTGAN import DataTransformer, BaseSynthesiser
from .utils import GeneralTransformer, DPSynthesiser
import opacus
from opacus import autograd_grad_sample, PrivacyEngine, utils
import dill
class Encoder(Module):
def __init__(self, data_dim, compress_dims, embedding_dim):
super(Encoder, self).__init__()
dim = data_dim
seq = []
for item in list(compress_dims):
seq += [Linear(dim, item), ReLU()]
dim = item
self.seq = Sequential(*seq)
self.fc1 = Linear(dim, embedding_dim)
self.fc2 = Linear(dim, embedding_dim)
def forward(self, input):
feature = self.seq(input)
mu = self.fc1(feature)
logvar = self.fc2(feature)
std = torch.exp(0.5 * logvar)
return mu, std, logvar
class Decoder(Module):
def __init__(self, embedding_dim, decompress_dims, data_dim):
super(Decoder, self).__init__()
dim = embedding_dim
seq = []
for item in list(decompress_dims):
seq += [Linear(dim, item), ReLU()]
dim = item
seq.append(Linear(dim, data_dim))
self.seq = Sequential(*seq)
self.sigma = Parameter(torch.ones(data_dim) * 0.1)
def forward(self, input):
return self.seq(input), self.sigma
def loss_function(recon_x, x, sigmas, mu, logvar, output_info, factor):
st = 0
loss = []
for column_info in output_info:
for span_info in column_info:
if len(column_info) != 1 or span_info.activation_fn != "softmax":
ed = st + span_info.dim
std = sigmas[st]
loss.append(((x[:, st] - torch.tanh(recon_x[:, st])) ** 2 / 2 / (std ** 2)).sum())
loss.append(torch.log(std) * x.size()[0])
st = ed
else:
ed = st + span_info.dim
loss.append(cross_entropy(recon_x[:, st:ed], torch.argmax(x[:, st:ed], dim=-1), reduction="sum"))
st = ed
assert st == recon_x.size()[1]
KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
return sum(loss) * factor / x.size()[0], KLD / x.size()[0]
class DPTVAE(BaseSynthesiser):
def __init__(
self,
embedding_dim=128,
compress_dims=(128, 128),
decompress_dims=(128, 128),
l2scale=1e-5,
batch_size=500,
disabled_dp=False,
delta=1e-5,
noise_multiplier=3.5,
max_per_sample_grad_norm=1.0,
epsilon=1.0,
iterations=300,
verbose=True,
):
self.embedding_dim = embedding_dim
self.compress_dims = compress_dims
self.decompress_dims = decompress_dims
self.l2scale = l2scale
self.batch_size = batch_size
self.loss_factor = 2
self.iterations = iterations
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# opacus parameters
self.noise_multiplier = noise_multiplier
self.disabled_dp = disabled_dp
self.delta = delta
self.max_per_sample_grad_norm = max_per_sample_grad_norm
self.epsilon = epsilon
self.epsilon_list = []
self.alpha_list = []
self.loss_list = []
self.verbose = verbose
def train(self, data, categorical_columns=None, ordinal_columns=None, update_epsilon=None, verbose=False):
self.transformer = DataTransformer()
self.transformer.fit(data, discrete_columns=categorical_columns)
data = self.transformer.transform(data)
dataset = TensorDataset(torch.from_numpy(data.astype("float32")).to(self.device))
loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True, drop_last=True)
data_dim = self.transformer.output_dimensions
self.encoder = Encoder(data_dim, self.compress_dims, self.embedding_dim).to(self.device)
self.decoder = Decoder(self.embedding_dim, self.compress_dims, data_dim).to(self.device)
self.optimizerAE = Adam(list(self.encoder.parameters()) + list(self.decoder.parameters()), weight_decay=self.l2scale)
privacy_engine = opacus.PrivacyEngine(
self.decoder,
batch_size=self.batch_size,
sample_size=data.shape[0],
alphas=[1 + x / 10.0 for x in range(1, 100)] + list(range(12, 64)),
target_delta=self.delta,
noise_multiplier=self.noise_multiplier,
max_grad_norm=self.max_per_sample_grad_norm,
clip_per_layer=True,
)
if not self.disabled_dp:
privacy_engine.attach(self.optimizerAE)
if hasattr(self, "privacy_engine"):
epsilon, best_alpha = self.optimizerAE.privacy_engine.get_privacy_spent(self.delta)
else:
epsilon = 0
for i in range(self.iterations):
if not self.disabled_dp:
if self.epsilon < epsilon:
break
for id_, data_ in enumerate(loader):
self.optimizerAE.zero_grad()
real = data_[0].to(self.device)
mu, std, logvar = self.encoder(real)
eps = torch.randn_like(std)
emb = eps * std + mu
rec, sigmas = self.decoder(emb)
loss_1, loss_2 = loss_function(rec, real, sigmas, mu, logvar, self.transformer.output_info_list, self.loss_factor)
loss = loss_1 + loss_2
loss.backward()
self.optimizerAE.step()
self.decoder.sigma.data_.clamp_(0.01, 1.0)
if not self.disabled_dp:
for p in self.decoder.parameters():
if hasattr(p, "grad_sample"):
del p.grad_sample
epsilon, best_alpha = self.optimizerAE.privacy_engine.get_privacy_spent(self.delta)
self.epsilon_list.append(epsilon)
self.alpha_list.append(best_alpha)
if verbose:
print("eps: {:f} \t alpha: {:f} \t Loss: {:f}".format(epsilon, best_alpha, loss.detach().cpu()))
if not self.disabled_dp:
if self.epsilon < epsilon:
break
self.loss_list.append(loss)
privacy_engine.detach()
self.privacy_engine = privacy_engine
self.state_dict = self.optimizerAE.state_dict()
return self.loss_list, self.epsilon_list, self.alpha_list
def sample(self, samples):
self.decoder.eval()
steps = samples // self.batch_size + 1
data_ = []
for _ in range(steps):
mean = torch.zeros(self.batch_size, self.embedding_dim)
std = mean + 1
noise = torch.normal(mean=mean, std=std).to(self.device)
fake, sigmas = self.decoder(noise)
fake = torch.tanh(fake)
data_.append(fake.detach().cpu().numpy())
data_ = np.concatenate(data_, axis=0)
data_ = data_[:samples]
return self.transformer.inverse_transform(data_, sigmas.detach().cpu().numpy())
def set_device(self, device):
self.device = device
self.decoder.to(self.device)
def save(self, path):
assert hasattr(self, "data_sampler")
# always save a cpu model.
device_bak = self.device
self.device = torch.device("cpu")
self.encoder.to(self.device)
self.decoder.to(self.device)
torch.save(self, path, pickle_module=dill)
self.device = device_bak
self.encoder.to(self.device)
self.decoder.to(self.device)
@classmethod
def load(cls, path):
model = torch.load(path)
model.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.encoder.to(model.device)
model.decoder.to(model.device)
return model
def DPTVAE_runner(seed, run_name, epsilon, delta, n_s, df, batch_size, save, load, categorical_columns, ordinal_columns, iterations, noise_multiplier):
if not load:
gan = DPSynthesiser(
DPTVAE(
batch_size=batch_size,
iterations=iterations,
delta=delta,
noise_multiplier=noise_multiplier,
),
GeneralTransformer(),
epsilon=epsilon,
)
gan.fit(df, categorical_columns=categorical_columns, ordinal_columns=ordinal_columns, verbose=True, seed=seed)
else:
gan = DPTVAE.load(load)
df_synth = gan.sample(n_s)
if save:
gan.save(save)
return df_synth, None
Any help would be appreciated, thanks!
