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