def _sampler(self, samples=1000):
d_ = torch.ones(samples)
if d == 1:
# If SZ is adopted, then some Districts and Schools buy in
dist = Poisson(self.n_districts)\
.sample([samples])\
.reshape([samples])
schools = NegativeBinomial(tensor([3.]),
tensor([0.8]))\
.sample([samples, self.n_districts.int()])\
.sum(dim=1)\
.reshape([samples])
sz = 15000. * dist + 2430 * schools
else:
dist, schools, sz = torch.zeros(samples),\
torch.zeros(samples),\
torch.zeros(samples)
if d < 2:
sf = LogNormal(
*self._lognormal_params(300000., 10000.))\
.sample([samples])
else:
sf = torch.zeros(samples)
# System & Infrastructure
az = LogNormal(self.az_means[d], self.az_sds[d]).sample([samples])
salary_estimate = Normal(70000., 5000.).sample([samples])
fa = Beta(self.fa_ms[d], self.fa_ks[d]).sample([samples])
dt = Beta(self.dt_ms[d], self.dt_ks[d]).sample([samples])
return d_, dist, schools, sz, az, sf, fa, dt
def gen_bid_ask(x0,
x,
vol=0.15,
spread_mult=0.8,
stoch_spread_mult=0.5,
clampSpread=True,
vol_window=10):
"""
Generates bid/ask based on the rolling vol of the input mid process and stochastic spread add-on
Stochastic spread is a function of mid process volatility
:param vol: volatility of mid process
:param spread_mult: multiplier of resulting spread
:param stoch_spread_mult: multiplier of sochastic addition on the top of rolling vol multiple
:param clampSpread: determines whether to imply bounds on the resulting spread after stochastic lognormal addition is applied
"""
nsteps = len(x)
rolling_vol = rolling_stdev(x, vol_window) / math.sqrt(nsteps)
spread_bid = rolling_vol + LogNormal(
0, stoch_spread_mult * x0 * vol / math.sqrt(nsteps)).sample((nsteps, ))
spread_ask = rolling_vol + LogNormal(
0, stoch_spread_mult * x0 * vol / math.sqrt(nsteps)).sample((nsteps, ))
# clamp resulting spread to guard against unreasonable samples from lognormal distr
if (clampSpread):
mean_vol = torch.mean(rolling_vol)
torch.clamp(spread_bid, 0.1 * mean_vol, 2.5 * mean_vol)
torch.clamp(spread_ask, 0.1 * mean_vol, 2.5 * mean_vol)
x_bid = x - spread_mult * spread_bid
x_ask = x + spread_mult * spread_ask
return x_ask, x_bid
def forward(self, x):
concentration = softplus(self.dir_concentration)
loc = self.dir_loc / self.dir_loc .norm(dim=-1, keepdim=True)
self.dir_sampler = PowerSpherical(loc, concentration)
e = self.nonlinear(self.embed1(x))
e = self.nonlinear(self.embed2(e))
self.rad_sampler = LogNormal(self.rad_mu(e), softplus(self.rad_scale(e)))
self.rad_sampler1 = LogNormal(self.rad_mu1(e), softplus(self.rad_scale1(e)))
self.bias_sampler = Normal(self.bias_mu, softplus(self.bias_scale))
direction_sample = self.dir_sampler.rsample()
radius_sample = self.rad_sampler.rsample()
radius_sample = (radius_sample * self.rad_sampler1.rsample()) ** 0.5
radius_sample = radius_sample ** 0.5
# radius_sample = LogNormal(self.rad_mu, softplus(self.rad_scale)).rsample()
# radius_sample = (radius_sample * LogNormal(self.rad_mu1, softplus(self.rad_scale1)).rsample()) ** 0.5
bias = self.bias_sampler.rsample() if self.bias else None
# weight = direction_sample * radius_sample.unsqueeze(0) ** 0.5
weight = direction_sample
output = F.linear(x*radius_sample, weight, bias)
return output
def __init__(self, loc, scale, validate_args=None, transform=None):
TModule.__init__(self)
LogNormal.__init__(self,
loc=loc,
scale=scale,
validate_args=validate_args)
self._transform = transform
def __init__(self):
super(Model, self).__init__()
self.fc1 = nn.Linear(28 * 28, 400)
# nn.init.uniform_(self.fc1.bias, -0.7, -0.7)
# nn.init.normal_(self.fc1.bias, -100, 99.3)
self.fc2 = nn.Linear(400, 400)
# nn.init.uniform_(self.fc2.bias, -0.7, -0.7)
# nn.init.normal_(self.fc2.bias, -100, 99.3)
self.fc3 = nn.Linear(400, 10)
self.activate = nn.ReLU()
self.name = 'mlp'
ln = LogNormal(0, 1)
self.mask1 = nn.Parameter(ln.sample(torch.Size([400])) - 1.5)
# nn.init.uniform_(self.mask1, -1, 0)
self.mask2 = nn.Parameter(ln.sample(torch.Size([400])) - 1.5)
def __init__(self, nsamples, multi):
self.npars = 3
self.nsamples = nsamples
self.multi = multi
loc = torch.tensor(np.log((0.1, 1., 20.)), dtype=torch.float32)
scale = torch.tensor(np.sqrt((0.05, 0.05, 0.05)), dtype=torch.float32)
self.prior = LogNormal(loc, scale)
def forward(self, inputs):
h1 = F.softplus(self.fc1(inputs))
h2 = F.softplus(self.fc2(h1))
mu = self.bnmu(self.fcmu(h2))
lv = self.bnlv(self.fclv(h2))
dist = LogNormal(mu, (0.5 * lv).exp())
return dist
def forward(self, input, sample=False):
# self.dir_loc.data /= torch.sum(self.dir_loc.data ** 2, dim=-1, keepdim=True) ** 0.5
# direction_sample = self.dir_rsampler(1, sample)[0]
if sample:
direction_sample = PowerSpherical(
self.dir_loc,
softplus(self.dir_softplus_inv_concentration)).rsample()
radius_sample = LogNormal(self.rad_mu,
softplus(self.rad_rho)).rsample()
else:
direction_sample = PowerSpherical(
self.dir_loc,
softplus(self.dir_softplus_inv_concentration)).mean
radius_sample = LogNormal(self.rad_mu, softplus(self.rad_rho)).mean
weight = direction_sample * radius_sample #.unsqueeze(-1)
return F.linear(input, weight, self.bias)
def test_Inference(self):
# ===== Distributions ===== #
dist = Normal(0., 1.)
mvn = Independent(Normal(torch.zeros(2), torch.ones(2)), 1)
# ===== Define model ===== #
linear = AffineProcess((f, g), (1., 0.25), dist, dist)
model = LinearGaussianObservations(linear, scale=0.1)
mv_linear = AffineProcess((fmvn, gmvn), (0.5, 0.25), mvn, mvn)
mvnmodel = LinearGaussianObservations(mv_linear, torch.eye(2), scale=0.1)
# ===== Test for multiple models ===== #
priors = Exponential(1.), LogNormal(0., 1.)
hidden1d = AffineProcess((f, g), priors, dist, dist)
oned = LinearGaussianObservations(hidden1d, 1., scale=0.1)
hidden2d = AffineProcess((fmvn, gmvn), priors, mvn, mvn)
twod = LinearGaussianObservations(hidden2d, torch.eye(2), scale=0.1 * torch.ones(2))
particles = 1000
# ====== Run inference ===== #
for trumod, model in [(model, oned), (mvnmodel, twod)]:
x, y = trumod.sample_path(1000)
algs = [
(NESS, {'particles': particles, 'filter_': APF(model.copy(), 200)}),
(NESS, {'particles': particles, 'filter_': UKF(model.copy())}),
(SMC2, {'particles': particles, 'filter_': APF(model.copy(), 200)}),
(SMC2FW, {'particles': particles, 'filter_': APF(model.copy(), 200)}),
(NESSMC2, {'particles': particles, 'filter_': APF(model.copy(), 200)})
]
for alg, props in algs:
alg = alg(**props).initialize()
alg = alg.fit(y)
w = normalize(alg._w_rec if hasattr(alg, '_w_rec') else torch.ones(particles))
tru_params = trumod.hidden.theta._cont + trumod.observable.theta._cont
inf_params = alg.filter.ssm.hidden.theta._cont + alg.filter.ssm.observable.theta._cont
for trup, p in zip(tru_params, inf_params):
if not p.trainable:
continue
kde = p.get_kde(weights=w)
transed = p.bijection.inv(trup)
densval = kde.logpdf(transed.numpy().reshape(-1, 1))
priorval = p.distr.log_prob(trup)
assert (densval > priorval.numpy()).all()
def construct_variational_parameters(T, N):
# T-1 params for the second part of variational Beta factors
kappa = Variable(Uniform(0, 2).rsample([T - 1]), requires_grad=True)
# T scale params for the variational Gamma factors
tau_0 = Uniform(0, 100).rsample([T])
# T rate params for the variational Gamma factors
tau_1 = LogNormal(0, 1).rsample([T])
tau = Variable(torch.stack((tau_0, tau_1)).T, requires_grad=True)
phi = Variable(
Dirichlet(1 / T * torch.ones(T)).rsample([N]),
requires_grad=True) # N,T params for the variational Cat factors
return kappa, tau, phi
def recon_model_loglik(self, x, eta):
# WARNING : the gaussian likelidhood is not supported
if self.likelihood == 'gaussian':
x_in = self.Psi.t() @ torch.log(x + 1).t()
diff = (x - eta)**2
sigma_sq = torch.exp(self.log_sigma_sq)
# No dimension constant as we sum after
return 0.5 * (-diff / sigma_sq - LOG_2_PI - self.log_sigma_sq)
elif self.likelihood == 'multinomial':
logp = (self.Psi.t() @ eta.t()).t()
mult_loss = Multinomial(logits=logp).log_prob(x).mean()
return mult_loss
elif self.likelihood == 'lognormal':
logp = F.logsoftmax((self.Psi.t() @ eta.t()).t(), axis=-1)
logN = torch.log(x.sum(axis=-1))
mu = logp + logN
sigma_sq = torch.exp(self.log_sigma_sq)
nz = x > 0
logn_loss = LogNormal(loc=mu[nz], scale=sigma_sq).log_prob(x[nz])
return logn_loss.mean()
else:
raise ValueError(
f'{self.likelihood} has not be properly specified.')
def __init__(self, kappa, gamma, sigma, **kwargs):
"""
Defines a Verhulst process.
:param kappa: The reversion parameter
:param gamma: The mean parameter
:param sigma: The standard deviation
"""
def f(x, k, g, s):
return k * (g - x) * x
def g_(x, k, g, s):
return s * x
super().__init__((f, g_), (kappa, gamma, sigma), LogNormal(0., 1.),
Normal(0., 1.), **kwargs)
def test_lnσ_transformation(self):
model = UnivariateGaussian()
two = torch.tensor(2.)
self.assertClose(model.σ_to_η.inv(two), torch.exp(two))
self.assertClose(model.σ_to_η(two), torch.log(two))
σ_prior, η_prior = LogNormal(0, 10), model.η_prior
# η_prior should just be N(0,10)
ηs = torch.randn(50)
self.assertClose(η_prior.log_prob(ηs), Normal(0, 10).log_prob(ηs))
# Transformed log density should include the log abs determinant of the
# inverse transform σ = log(η), which is log(η)
σs = model.σ_to_η.inv(ηs)
self.assertClose(η_prior.log_prob(ηs),
σ_prior.log_prob(σs) + torch.log(σs))
def model(self, input, targets):
priors = {}
for i, dense in enumerate(self.dense):
dw_prior = Normal(loc=torch.zeros_like(dense.weight),
scale=torch.ones_like(dense.weight))
db_prior = Normal(loc=torch.zeros_like(dense.bias),
scale=torch.ones_like(dense.bias))
priors['dense[{}].weight'.format(i)] = dw_prior
priors['dense[{}].bias'.format(i)] = db_prior
# set prior on scale
scale = LogNormal(0.0, 1.0)
# lift module parameters to random variables sampled from the priors
lifted_module = pyro.random_module("bnn", self, priors)
# sample a regressor (which also samples w and b)
lifted_reg_model = lifted_module()
pred = lifted_reg_model(input)
pyro.sample("obs", Normal(loc=pred, scale=scale), obs=targets)
from time import time torch.manual_seed(0) def get_x(theta, design): theta1 = theta[:, 0:1] theta2 = theta[:, 1:2] theta3 = theta[:, 2:3] x = 400. * theta2 * (torch.exp(-theta1*design) - torch.exp(-theta2*design)) / (theta3*(theta2-theta1)) return x n_inner = 1000 n_outer = 100 loc = torch.tensor(np.log((0.1, 1., 20.)), dtype=torch.float64) scale = torch.tensor(np.sqrt((0.05, 0.05, 0.05)), dtype=torch.float64) prior = LogNormal(loc, scale) prior = Independent(prior, 1) theta_inner = prior.sample((n_inner,)) theta_outer = prior.sample((n_outer,)) loc = torch.zeros(15, dtype=torch.float64) scale = 0.1 * torch.ones(15, dtype=torch.float64) noise = Normal(loc, scale) noise = Independent(noise, 1) noise_entropy = noise.entropy() noise_outer = noise.sample((n_outer,)) def objective(design): x_outer = get_x(theta_outer, design) x_inner = get_x(theta_inner, design) y_outer = x_outer + noise_outer # Get matrix of all y_outer-x_inner values
with open('outputs/pk_gda_K1.pkl', 'rb') as infile:
out_GDA = pickle.load(infile)
with open('outputs/pk_sgd.pkl', 'rb') as infile:
out_SGD = pickle.load(infile)
design_adv = torch.tensor(out_GDA['final_design'][0,:])
design_fig = torch.tensor(out_SGD['final_design'][0,:])
#######################
## PRIOR AND MODEL
#######################
loc = torch.tensor(np.log((0.1, 1., 20.)))
scale = torch.tensor(np.sqrt((0.05, 0.05, 0.05)))
prior = LogNormal(loc, scale)
prior = Independent(prior, 1)
def get_x(theta, design):
theta1 = theta[..., 0:1]
theta2 = theta[..., 1:2]
theta3 = theta[..., 2:3]
while design.dim() < theta1.dim():
design = design.unsqueeze(0)
x = 400. * theta2 * \
(torch.exp(-theta1*design) - torch.exp(-theta2*design)) \
/ (theta3*(theta2-theta1))
return x
loc = torch.zeros(design_adv.shape)
scale = 0.1 * torch.ones(design_adv.shape)
def update_rates(self, value):
rates = LogNormal(-value * value / 2.0, value).icdf(self.quantile)
self._rates = rates / (rates.sum() * self.probs)
if self._mu is not None:
self._rates *= self._mu.tensor
def standard_prior_like(posterior):
loc = torch.zeros_like(posterior.loc)
scale = torch.ones_like(posterior.scale)
prior = LogNormal(loc, scale)
return prior
IndependentTransform,
SigmoidTransform,
)
from sbi.utils import BoxUniform, MultipleIndependent, mcmc_transform, process_prior
from tests.user_input_checks_test import UserNumpyUniform
@pytest.mark.parametrize(
"prior, target_transform",
(
(Uniform(-torch.ones(1), torch.ones(1)), SigmoidTransform),
(BoxUniform(-torch.ones(2), torch.ones(2)), SigmoidTransform),
(UserNumpyUniform(torch.zeros(2), torch.ones(2)), SigmoidTransform),
(MultivariateNormal(torch.zeros(2), torch.eye(2)), AffineTransform),
(LogNormal(loc=torch.zeros(1), scale=torch.ones(1)), ExpTransform),
),
)
def test_transforms(prior, target_transform):
if isinstance(prior, UserNumpyUniform):
prior, *_ = process_prior(
prior,
dict(lower_bound=torch.zeros(2), upper_bound=torch.ones(2)),
)
transform = mcmc_transform(prior)
core_transform = transform._inv
if isinstance(core_transform, IndependentTransform):
core_transform = core_transform.base_transform
def kl_divergence(self, saver_net, trainer_net):
kld = 0
prev_weight_strength = nn.Parameter(
torch.Tensor(28 * 28, 1).uniform_(0, 0)).cuda()
alpha = self.alpha
if self.saved:
alpha = 1
for (saver_name,
saver_layer), (trainer_name,
trainer_layer) in zip(saver_net.items(),
trainer_net.items()):
trainer_dir_loc = trainer_layer['dir_loc']
trainer_dir_concentration = F.softplus(
trainer_layer['dir_softplus_inv_concentration'])
trainer_rad_mu = trainer_layer['rad_mu']
trainer_rad_sigma = F.softplus(trainer_layer['rad_rho'])
trainer_bias = trainer_layer['bias']
saver_dir_loc = saver_layer['dir_loc']
saver_dir_concentration = F.softplus(
saver_layer['dir_softplus_inv_concentration'])
saver_rad_mu = saver_layer['rad_mu']
saver_rad_sigma = F.softplus(saver_layer['rad_rho'])
saver_bias = saver_layer['bias']
fan_in, fan_out = _calculate_fan_in_and_fan_out(trainer_dir_loc)
concentration_init = ml_kappa(dim=fan_in, eps=self.model.eps)
if 'fc' in trainer_name:
std_init = math.sqrt((2 / fan_in) * self.model.ratio)
if 'conv' in trainer_name:
std_init = math.sqrt((2 / fan_out) * self.model.ratio)
out_features, in_features = saver_dir_loc.shape
saver_weight_strength = (std_init / saver_rad_sigma)
curr_strength = saver_weight_strength.expand(
out_features, in_features)
prev_strength = prev_weight_strength.permute(1, 0).expand(
out_features, in_features)
L2_strength = torch.max(curr_strength, prev_strength)
prev_weight_strength = saver_weight_strength
dir_loc_reg = (
(L2_strength * trainer_dir_loc * saver_dir_loc) /
(trainer_dir_loc.norm(2, dim=-1) *
saver_dir_loc.norm(2, dim=-1)).unsqueeze(-1)).sum()
q_dir = PowerSpherical(trainer_dir_loc, trainer_dir_concentration)
p_dir = PowerSpherical(saver_dir_loc, saver_dir_concentration)
kld_dir = KL_Powerspherical(q_dir, p_dir)
q_rad = LogNormal(trainer_rad_mu, trainer_rad_sigma)
p_rad = LogNormal(saver_rad_mu, saver_rad_sigma)
kld_rad = kl_divergence(q_rad, p_rad)
mu_bias_reg = ((trainer_bias - saver_bias) /
saver_rad_sigma.squeeze()).norm(2)**2
kld += kld_dir.sum(
) + 100 * kld_rad.sum() + 100 * mu_bias_reg + 100 * dir_loc_reg
return kld
def dists(self):
'''Returns the parametrized distributions for m1/m2.'''
# Creating the distributions always on the fly, otherwise we get
# PyTorch warnings about differentiating a second time.
return (LogNormal(self.m1m2_mean[0], torch.exp(self.m1m2_log_std[0])),
LogNormal(self.m1m2_mean[1], torch.exp(self.m1m2_log_std[1])))
print("J estimates: adv {:.1e}, unif {:.1e}".format(diagJ_adv, diagJ_uni))
diagSIG_adv = objective(design_adv)
diagSIG_uni = objective(design_uni)
print("SIG estimates: adv {:.1f}, unif {:.1f}".format(diagSIG_adv,
diagSIG_uni))
###############################
## POSTERIOR PLOT VS ADV DESIGN
###############################
loc = torch.tensor(np.log((0.1, 1., 20.)))
scale = torch.tensor(np.sqrt((0.05, 0.05, 0.05)))
prior = LogNormal(loc, scale)
prior = Independent(prior, 1)
def get_x(theta, design):
theta1 = theta[..., 0:1]
theta2 = theta[..., 1:2]
theta3 = theta[..., 2:3]
while design.dim() < theta1.dim():
design = design.unsqueeze(0)
x = 400. * theta2 * \
(torch.exp(-theta1*design) - torch.exp(-theta2*design)) \
/ (theta3*(theta2-theta1))
return x
class NewLinear(nn.Module):
"""docstring for NewLinear"""
def __init__(self, in_features, out_features, bias=True, noise_shape=1):
super(NewLinear, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.bias = bias
self.dir_concentration = nn.Parameter(torch.Tensor(out_features))
self.dir_loc = nn.Parameter(torch.Tensor(out_features, in_features))
nn.init.kaiming_normal_(self.dir_loc)
nn.init.normal_(self.dir_concentration, out_features*10, 1)
self.rad_mu = nn.Linear(in_features, in_features)
self.rad_scale = nn.Linear(in_features, in_features)
self.rad_mu1 = nn.Linear(in_features, in_features)
self.rad_scale1 = nn.Linear(in_features, in_features)
self.embed1 = nn.Linear(in_features, in_features)
self.embed2 = nn.Linear(in_features, in_features)
self.nonlinear = nn.ReLU()
# self.rad_mu = nn.Parameter(torch.Tensor(in_features))
# self.rad_scale = nn.Parameter(torch.Tensor(in_features))
# self.rad_mu1 = nn.Parameter(torch.Tensor(in_features))
# self.rad_scale1 = nn.Parameter(torch.Tensor(in_features))
# nn.init.normal_(self.rad_mu, math.log(2.0), 0.0001)
# nn.init.normal_(self.rad_scale, softplus_inv(0.0001), 0.0001)
# nn.init.normal_(self.rad_mu1, math.log(2.0), 0.0001)
# nn.init.normal_(self.rad_scale1, softplus_inv(0.0001), 0.0001)
self.bias_mu = nn.Parameter(torch.Tensor(out_features))
self.bias_scale = nn.Parameter(torch.Tensor(out_features))
nn.init.normal_(self.bias_mu, 0.0, 0.0001)
nn.init.normal_(self.bias_scale, softplus_inv(0.0001), 0.0001)
def forward(self, x):
concentration = softplus(self.dir_concentration)
loc = self.dir_loc / self.dir_loc .norm(dim=-1, keepdim=True)
self.dir_sampler = PowerSpherical(loc, concentration)
e = self.nonlinear(self.embed1(x))
e = self.nonlinear(self.embed2(e))
self.rad_sampler = LogNormal(self.rad_mu(e), softplus(self.rad_scale(e)))
self.rad_sampler1 = LogNormal(self.rad_mu1(e), softplus(self.rad_scale1(e)))
self.bias_sampler = Normal(self.bias_mu, softplus(self.bias_scale))
direction_sample = self.dir_sampler.rsample()
radius_sample = self.rad_sampler.rsample()
radius_sample = (radius_sample * self.rad_sampler1.rsample()) ** 0.5
radius_sample = radius_sample ** 0.5
# radius_sample = LogNormal(self.rad_mu, softplus(self.rad_scale)).rsample()
# radius_sample = (radius_sample * LogNormal(self.rad_mu1, softplus(self.rad_scale1)).rsample()) ** 0.5
bias = self.bias_sampler.rsample() if self.bias else None
# weight = direction_sample * radius_sample.unsqueeze(0) ** 0.5
weight = direction_sample
output = F.linear(x*radius_sample, weight, bias)
return output
def kl_divergence(self):
pass
def expand(self, batch_shape):
return LogNormal.expand(self, batch_shape, _instance=self)
def forward(self, hidden):
mu = self.bnmu(self.fcmu(hidden))
lv = self.bnlv(self.fclv(hidden))
dist = LogNormal(mu, (0.5 * lv).exp())
return dist
def make(self,
mode: int,
num_clients: int,
show_plots: bool = False,
**kwargs) -> None:
if os.path.exists(self.root_dir / "client_data"):
shutil.rmtree(self.root_dir / "client_data")
client_data_path = Path(self.root_dir / "client_data")
client_data_path.mkdir()
if not isinstance(self.test_data.targets, torch.Tensor):
self.test_data.targets = torch.tensor(self.test_data.targets)
test_data = [self.test_data[j] for j in range(len(self.test_data))]
torch.save(test_data, client_data_path / "test_data.pth")
if mode == 0: # IID
# Shuffle data
data_ids = torch.randperm(self.num_train_data, dtype=torch.int32)
num_data_per_client = self.num_train_data // num_clients
if not isinstance(self.train_data.targets, torch.Tensor):
self.train_data.targets = torch.tensor(self.train_data.targets)
pbar = tqdm(range(num_clients), desc=f"{self.dataset_name} IID: ")
for i in pbar:
client_path = Path(client_data_path / str(i))
client_path.mkdir()
# TODO: Make this parallel for large number of clients & large datasets (Maybe not required)
train_data = [
self.train_data[j]
for j in data_ids[i * num_data_per_client:(i + 1) *
num_data_per_client]
]
pbar.set_postfix({'# data / Client': num_data_per_client})
if show_plots:
self._plot(train_data,
title=f"Client {i+1} Data Distribution")
# Split data equally and send to the client
torch.save(train_data, client_data_path / str(i) / "data.pth")
elif mode == 1: # Non IID Balanced
num_data_per_client = self.num_train_data // num_clients
classs_sampler = Dirichlet(
torch.empty(self.num_classes).fill_(kwargs.get('dir_alpha')))
# print(torch.empty(self.num_classes).fill_(2.0))
if not isinstance(self.train_data.targets, torch.Tensor):
self.train_data.targets = torch.tensor(self.train_data.targets)
assigned_ids = []
pbar = tqdm(range(num_clients),
desc=f"{self.dataset_name} Non-IID Balanced: ")
for i in pbar:
client_path = Path(client_data_path / str(i))
client_path.mkdir()
# Compute class prior probabilities for each client
p_ij = classs_sampler.sample(
) # Share of jth class for ith client (always sums to 1)
# print(p_ij)
weights = torch.zeros(self.num_train_data)
# print(torch.nonzero(self.train_data.targets == 9))
for c_id in range(self.num_classes):
weights[self.train_data.targets == c_id] = p_ij[c_id]
weights[
assigned_ids] = 0.0 # So that previously assigned data are not sampled again
# Sample each data point uniformly without replacement based on
# the sampling probability assigned based on its class
data_ids = torch.multinomial(weights,
num_data_per_client,
replacement=False)
train_data = [self.train_data[j] for j in data_ids]
# print(f"Client {i} has {len(train_data)} data points.")
pbar.set_postfix({'# data / Client': len(train_data)})
assigned_ids += data_ids.tolist()
torch.save(train_data, client_data_path / str(i) / "data.pth")
if show_plots:
self._plot(train_data,
title=f"Client {i+1} Data Distribution")
elif mode == 2: # Non IID Unbalanced
num_data_per_client = self.num_train_data // num_clients
num_data_per_class = self.num_train_data / (self.num_classes *
num_clients)
classs_sampler = Dirichlet(
torch.empty(self.num_classes).fill_(kwargs.get('dir_alpha')))
assigned_ids = []
pbar = tqdm(range(num_clients),
desc=f"{self.dataset_name} Non-IID Unbalanced: ")
if not isinstance(self.train_data.targets, torch.Tensor):
self.train_data.targets = torch.tensor(self.train_data.targets)
for i in pbar:
train_data = []
client_path = Path(client_data_path / str(i))
client_path.mkdir()
# Compute class prior probabilities for each client
p_ij = classs_sampler.sample(
) # Share of jth class for ith client (always sums to 1)
c_sampler = Categorical(p_ij)
data_sampler = LogNormal(
torch.tensor(num_data_per_class).log(),
kwargs.get('lognorm_std'))
while (True):
num_data_left = num_data_per_client - len(train_data)
c = c_sampler.sample()
num_data_c = int(data_sampler.sample())
# print(c, num_data_c, len(train_data))
data_ids = torch.nonzero(
self.train_data.targets == c.item()).flatten()
# data_ids = [x for x in data_ids if x not in assigned_ids] # Remove duplicated ids
# print(data_ids.shape)
num_data_c = min(num_data_c, data_ids.shape[0])
if num_data_c >= num_data_left:
train_data += [
self.train_data[j]
for j in data_ids[:num_data_left]
]
break
else:
train_data += [
self.train_data[j] for j in data_ids[:num_data_c]
]
assigned_ids += data_ids[:num_data_c].tolist()
pbar.set_postfix({'# data / Client': len(train_data)})
torch.save(train_data, client_data_path / str(i) / "data.pth")
if show_plots:
self._plot(train_data,
title=f"Client {i+1} Data Distribution")
else:
raise ValueError("Unknown mode. Mode must be {0,1}")
def forward(self, hidden):
# hidden size: (batch_size x hidden_size)
mu = self.bnmu(self.fcmu(hidden))
lv = self.bnlv(self.fclv(hidden))
dist = LogNormal(mu, (0.5 * lv).exp())
return dist