def odin_infer(loader,
model,
num_bs,
num_classes,
with_acc=False,
seed=0,
T=1000,
eps=0.0014):
loss_fn = torch.nn.CrossEntropyLoss()
torch.manual_seed(seed)
model.eval()
a_test_ = Exponential(torch.ones([1, 400]))
a_test = a_test_.sample((num_bs, ))
acc = 0.
outputs = np.zeros([num_bs, len(loader.dataset), num_classes + 1])
beg = 0
for i, (img, label) in enumerate(loader):
index = list(range(beg, beg + img.size(0)))
beg = beg + img.size(0)
label = label.numpy().reshape(-1, 1)
img_ = img.cuda()
img_.requires_grad = True
output = model(img_, torch.zeros([img.shape[0], 400]))
output = output / T
pseudo_label = output.argmax(-1).cuda()
loss = loss_fn(output, pseudo_label)
loss.backward()
gradient = torch.ge(img_.grad.data, 0)
gradient = (gradient.float() - 0.5) * 2
gradient.index_copy_(
1,
torch.LongTensor([0]).cuda(),
gradient.index_select(1,
torch.LongTensor([0]).cuda()) / (0.2023))
gradient.index_copy_(
1,
torch.LongTensor([1]).cuda(),
gradient.index_select(1,
torch.LongTensor([1]).cuda()) / (0.1994))
gradient.index_copy_(
1,
torch.LongTensor([2]).cuda(),
gradient.index_select(1,
torch.LongTensor([2]).cuda()) / (0.2010))
img_new = torch.add(img_.data, -eps, gradient)
for _ in range(num_bs):
w_test = a_test[_].repeat_interleave(img.shape[0], dim=0)
output_new = model(img_new, w_test).cpu().detach().numpy()
outputs[_, index] = np.concatenate([output_new, label], axis=1)
if with_acc:
pred = outputs.sum(0)[:, :-1].argmax(1)
label = outputs[0][:, -1]
acc = pred == label
print(f"[Test] acc : {acc.mean()}")
return outputs
def forward(self, x, t, mask):
#t : TxBSx2
d_j = t[:, :, 0] #TxBS
batch_size, seq_length = x.size(1), x.size(0)
past_influences = []
for idx in range(self.m):
t_pad = torch.cat(
[torch.zeros(idx + 1, batch_size).to(device),
d_j])[:-(idx + 1), :][:, :, None] #TxBSx1
past_influences.append(t_pad)
past_influences = torch.cat(
past_influences, dim=-1) * self.alpha[None, None, :] #TxBSxm
total_influence = torch.sum(past_influences,
dim=-1) + self.gamma[None, :].exp()
#To consider from time step 1
m = Exponential(total_influence[1:, :]) #T-1xBS
ll_loss = (m.log_prob(d_j[1:, :])).sum()
metric_dict = {
'true_ll': -ll_loss.detach(),
"marker_acc": 0.,
"marker_acc_count": 1.
}
with torch.no_grad():
time_mse = ((d_j[1:, :] - 1. / total_influence[1:, :]) *
mask[1:, :])**2.
metric_dict['time_mse'] = time_mse.sum().detach().cpu().numpy()
metric_dict['time_mse_count'] = mask[
1:, :].sum().detach().cpu().numpy()
return -ll_loss, metric_dict
def __init__(self, in_features, a=None, trainable=True):
'''
Initialization.
Args:
in_features: shape of the input
a: trainable parameter
trainable: sets `a` as a trainable parameter
`a` is initialized to 1 by default, higher values = higher-frequency,
5-50 is a good starting point if you already think your data is periodic,
consider starting lower e.g. 0.5 if you think not, but don't worry,
`a` will be trained along with the rest of your model
'''
super(Snake, self).__init__()
self.in_features = in_features if isinstance(in_features,
list) else [in_features]
# Initialize `a`
if a is not None:
self.a = Parameter(torch.ones(self.in_features) *
a) # create a tensor out of alpha
else:
m = Exponential(torch.tensor([0.1]))
self.a = Parameter(
(m.rsample(self.in_features)
).squeeze()) # random init = mix of frequencies
self.a.requiresGrad = trainable # set the training of `a` to true
def _train_a_query(self, indice):
indexer = np.stack(
[np.array(indice[i::self.num_groups])
for i in range(self.num_groups)])
sampler = SubsetRandomSampler(indice)
loader = DataLoader(self.dataset, batch_size=128, sampler=sampler,
num_workers=4, pin_memory=True)
alpha_generator = Exponential(torch.ones([1, self.num_groups]))
t_iter = tqdm(range(self.num_epoch),
total=self.num_epoch,
desc="Training")
for t in t_iter:
alpha = alpha_generator.sample().cuda()
for img, label, index in loader:
self.model.train()
n0 = img.size(0)
u_is = []
for i in index:
u_i = np.where(indexer == i.item())[0][0]
u_is += [u_i]
w = alpha[0, u_is].cuda()
output = self.model(img.cuda(), alpha.repeat_interleave(n0, 0))
loss = self.loss_fn(output, label.cuda(), w)
self.optim.zero_grad()
loss.backward()
self.optim.step()
self.sched.step()
def __init__(self, rate, upper):
self.base = Exponential(rate)
self._batch_shape = self.base.rate.size()
self._upper = upper
self.upper = torch.full_like(self.base.rate, upper)
# normaliser
self.normaliser = self.base.cdf(self.upper)
self.uniform = Uniform(torch.zeros_like(self.upper), self.normaliser)
def infer_a_sample(img, model, num_classes, num_bs, fac):
model.eval()
a_test = Exponential(torch.ones([1, 400]))
w_test = a_test.sample((num_bs, ))
output = np.zeros([num_bs, num_classes])
for _ in range(num_bs):
o = model(img, w_test[_], fac).cpu().numpy()
output[_] = o
return output
def __init__(self, config):
self.config = config
torch.manual_seed(config["random_seed"])
self.M = int(config["M"])
self.K = int(config["K"])
self.G = Exponential(1).sample((self.M, self.K)).float().cpu().detach().numpy()
self.G = self.G[:, np.argsort(self.G.sum(axis=0))]
self.P = float(config["transmission_power"])
def __init__(self, in_features, alpha=None, alpha_trainable=True):
"""
Init method.
"""
super(Snake, self).__init__()
self.in_features = in_features
if alpha is not None:
self.alpha = Parameter(torch.ones(in_features) * alpha)
else:
m = Exponential(torch.tensor([0.1]))
self.alpha = Parameter(m.sample(in_features))
self.alpha.requiresGrad = alpha_trainable
class ExpTimeToOpen(Base):
def __init__(self, rate):
self.dist = Exponential(rate)
def log_prob(self, times):
dt = times[:, 0]
return self.dist.log_prob(dt)
def __init__(self, scale, alpha, validate_args=None):
self.scale, self.alpha = broadcast_all(scale, alpha)
base_dist = Exponential(self.alpha, validate_args=validate_args)
transforms = [ExpTransform(), AffineTransform(loc=0, scale=self.scale)]
super(Pareto, self).__init__(base_dist,
transforms,
validate_args=validate_args)
class ExpTimeToOpen(Base):
def __init__(self, rate):
self.dist = Exponential(rate)
def log_prob(self, times):
dt = times[:, 1] - times[:, 0]
dt.apply_(lambda x: x + 24
if x < 0 else x) # find more time effective way to compute
return self.dist.log_prob(dt)
def __init__(self, scale, concentration, validate_args=None):
self.scale, self.concentration = broadcast_all(scale, concentration)
self.concentration_reciprocal = self.concentration.reciprocal()
base_dist = Exponential(torch.ones_like(self.scale), validate_args=validate_args)
transforms = [PowerTransform(exponent=self.concentration_reciprocal),
AffineTransform(loc=0, scale=self.scale)]
super(Weibull, self).__init__(base_dist,
transforms,
validate_args=validate_args)
def __init__(
self,
synapses,
neurons,
input_channel=1,
output_channel=1,
step=16,
leak=32,
fodep=None,
delay=None,
w_init=None,
theta=None,
dense=None,
):
super(RecurColumn, self).__init__()
self.synapses = synapses
self.neurons = neurons
self.input_channel = input_channel
self.output_channel = output_channel
assert theta or dense, 'either theta or dense should be specified'
self.theta = theta = theta or dense * (synapses * input_channel)
self.dense = dense = dense or theta / (synapses * input_channel)
assert dense < 2 * input_channel * \
synapses, 'invalid theta or density, try setting a smaller value'
w_init = w_init or dense
# default response function: StepFireLeak
self.response_function = StepFireLeak(step, leak)
self.fodep = fodep = fodep or self.response_function.fodep
assert fodep >= self.response_function.fodep, f'forced depression should be at least {self.response_function.fodep}'
self.delay = delay or self.fodep
# initialize weight to w_init
self.weight_i = nn.parameter.Parameter(Exponential(1 / w_init).sample(
(self.output_channel * self.neurons,
self.input_channel * self.synapses)).clip(0, 1),
requires_grad=True)
self.weight_s = nn.parameter.Parameter(Exponential(1 / w_init).sample(
(self.output_channel * self.neurons,
self.input_channel * self.synapses)).clip(0, 1),
requires_grad=True)
print(
f'Building full connected TNN layer with theta={theta:.4f}, dense={dense:.4f}, fodep={fodep}'
)
class HomogeneousPoissonProcess:
def __init__(self, rate=1):
self.rate = rate
self.exp = Exponential(rate)
def sample(self, size, max_seq_len, max_time=math.inf):
gaps = self.exp.sample((size, max_seq_len))
times = torch.cumsum(gaps, dim=1)
masks = (times <= max_time).float()
return times, masks
def _infer_gbs(self, with_acc=False, seed=0):
torch.manual_seed(seed)
a_test_ = Exponential(torch.ones([1, self.args.n_a]))
a_test = a_test_.sample([self.args.num_bs, ])
outputs = np.zeros([self.args.num_bs, len(self.loader.dataset), self.args.num_classes + 1])
beg = 0
for i, (img, label) in enumerate(self.loader):
end = beg + img.size(0)
label = label.numpy().reshape(-1, 1)
for _ in range(self.args.num_bs):
w_test = a_test[_].repeat_interleave(img.size(0), dim=0)
output = self._infer_a_batch_a_bs(img, w_test)
outputs[_, beg: end] = np.concatenate([output, label], axis=1)
beg = end
if with_acc:
pred = outputs.sum(0)[:, :-1].argmax(1)
label = outputs[0][:, -1]
acc = (pred == label).mean()
print(f"[Test] acc : {acc}")
return outputs
def odin_infer_a_sample(img,
model,
num_classes,
num_bs,
fac,
T=1000,
eps=0.0001):
model.eval()
loss_fn = torch.nn.CrossEntropyLoss()
img = img.cuda()
img.requires_grad = True
model.zero_grad()
a_test = Exponential(torch.ones([1, 400]))
w_test = a_test.sample((num_bs, ))
outputs = np.zeros([num_bs, num_classes])
for _ in range(num_bs):
# w_test = a_test[_].repeat_interleave(img.shape[0], dim=0)
img_ = img.cuda()
img_.requires_grad = True
output = model(img_, w_test[_], fac)
output = output / T
pseudo_label = output.argmax(-1).cuda()
loss = loss_fn(output, pseudo_label)
loss.backward()
gradient = torch.ge(img_.grad.data, 0)
# gradient = (gradient.float() - 0.5) * 2
# gradient.index_copy_(1, torch.LongTensor([0]).cuda(),
# gradient.index_select(1, torch.LongTensor([0]).cuda()) / (0.2023))
# gradient.index_copy_(1, torch.LongTensor([1]).cuda(),
# gradient.index_select(1, torch.LongTensor([1]).cuda()) / (0.1994))
# gradient.index_copy_(1, torch.LongTensor([2]).cuda(),
# gradient.index_select(1, torch.LongTensor([2]).cuda()) / (0.2010))
img_new = torch.add(img_.data, -eps, gradient)
output_new = model(img_new, w_test[_], fac).cpu().detach().numpy()
outputs[_] = output_new
return outputs
class RightTruncatedExponential(torch.distributions.Distribution):
def __init__(self, rate, upper):
self.base = Exponential(rate)
self._batch_shape = self.base.rate.size()
self._upper = upper
self.upper = torch.full_like(self.base.rate, upper)
# normaliser
self.normaliser = self.base.cdf(self.upper)
self.uniform = Uniform(torch.zeros_like(self.upper), self.normaliser)
def rsample(self, sample_shape=torch.Size()):
# sample from truncated support (0, normaliser)
# where normaliser = base.cdf(upper)
u = self.uniform.rsample(sample_shape)
x = self.base.icdf(u)
return x
def log_prob(self, value):
return self.base.log_prob(value) - torch.log(self.normaliser)
def cdf(self, value):
return self.base.cdf(value) / self.normaliser
def icdf(self, value):
return self.base.icdf(value * self.normaliser)
def cross_entropy(self, other):
assert isinstance(other, RightTruncatedExponential)
assert type(self.base) is type(
other.base) and self._upper == other._upper
a = torch.log(other.base.rate) - torch.log(other.normaliser)
log_b = torch.log(self.base.rate) + torch.log(
other.base.rate) - torch.log(self.normaliser)
b = torch.exp(log_b)
c = (torch.exp(-self.base.rate) *
(-self.base.rate - 1) + 1) / (self.base.rate**2)
return -(a - b * c)
def entropy(self):
return self.cross_entropy(self)
def infer(loader,
model,
num_bs,
num_classes,
with_acc=False,
with_indice=False,
is_mcd=False,
seed=0):
torch.manual_seed(seed)
model.eval()
a_test_ = Exponential(torch.ones([1, 400]))
a_test = a_test_.sample((num_bs, ))
acc = 0.
outputs = np.zeros([num_bs, len(loader.dataset), num_classes + 1])
beg = 0
indice = []
ret = []
for i, (img, label, _) in enumerate(loader):
img = img.cuda()
index = list(range(beg, beg + img.size(0)))
beg = beg + img.size(0)
label = label.numpy().reshape(-1, 1)
indice += [_]
for _ in range(num_bs):
w_test = a_test[_].repeat_interleave(img.shape[0], dim=0)
if is_mcd:
model.apply(apply_dropout)
output = model(img, w_test.cuda()).cpu().numpy()
outputs[_, index] = np.concatenate([output, label], axis=1)
ret += [outputs]
if with_acc:
pred = outputs.sum(0)[:, :-1].argmax(1)
label = outputs[0][:, -1]
acc = pred == label
print(f"[Test] acc : {acc.mean()}")
ret += [acc.mean()]
if with_indice:
ret += [torch.cat(indice)]
return ret
def __init__(self,
synapses,
neurons,
input_channel=1,
output_channel=1,
step=16,
leak=32,
bias=0.5,
winners=None,
fodep=None,
w_init=None,
theta=None,
dense=None,
alpha1=0.02,
alpha2=0.02,
beta1=0.99,
beta2=0.999):
super(FullColumn, self).__init__()
# model skeleton parameters
self.synapses = synapses
self.neurons = neurons
self.input_channel = input_channel
self.output_channel = output_channel
# threshold parameters
assert theta or dense, 'either theta or dense should be specified'
self.theta = theta = theta or dense * (synapses * input_channel)
self.dense = dense = dense or theta / (synapses * input_channel)
w_init = w_init or dense
# spiking control parameters
self.response_function = StepFireLeak(step, leak)
self.fodep = fodep = fodep or self.response_function.fodep
assert fodep >= self.response_function.fodep, f'forced depression should be at least {self.response_function.fodep}'
self.winners = winners = winners or neurons
# initialize weight and bias
self.bias = nn.parameter.Parameter(
torch.zeros(self.output_channel * self.neurons) + bias,
requires_grad=False)
self.weight = nn.parameter.Parameter(Exponential(1 / w_init).sample(
(self.output_channel * self.neurons,
self.input_channel * self.synapses)).clip(min=0),
requires_grad=True)
self.dual = SignalDualBackground()
self.optimizer = AdamBSTDP(self.weight, alpha1, alpha2, beta1, beta2)
print(
'Building full connected TNN layer with '
f'theta={theta:.4f}, '
f'dense={dense:.4f}, '
f'fodep={fodep}, ', f'winners={winners}, '
f'bias={bias}')
def __init__(self, *, structure: structure_prior.StructurePrior, intensity_mu_sig: tuple, lifetime: float,
frame_range: tuple, xy_unit: str, px_size: tuple, density=None, em_avg=None, intensity_th=None):
"""
Args:
structure:
intensity_mu_sig:
lifetime:
xy_unit:
px_size:
frame_range: specifies the frame range
density:
em_avg:
intensity_th:
"""
super().__init__(structure=structure,
photon_range=None,
xy_unit=xy_unit,
px_size=px_size,
density=density,
em_avg=em_avg)
self.n_sampler = np.random.poisson
self.frame_range = frame_range
self.intensity_mu_sig = intensity_mu_sig
self.intensity_dist = torch.distributions.normal.Normal(self.intensity_mu_sig[0],
self.intensity_mu_sig[1])
self.intensity_th = intensity_th if intensity_th is not None else 1e-8
self.lifetime_avg = lifetime
self.lifetime_dist = Exponential(1 / self.lifetime_avg) # parse the rate not the scale ...
self.t0_dist = torch.distributions.uniform.Uniform(*self._frame_range_plus)
"""
Determine the total number of emitters. Depends on lifetime, frames and emitters.
(lifetime + 1) because of binning effect.
"""
self._emitter_av_total = self._em_avg * self._num_frames_plus / (self.lifetime_avg + 1)
def __init__(self, in_features, alpha = None, alpha_trainable = True):
'''
Initialization.
INPUT:
- in_features: shape of the input
- alpha: trainable parameter
alpha is initialized to 1 by default, higher values = higher-frequency,
5-50 is a good starting point if you already think your data is periodic,
consider starting lower e.g. 0.5 if you think not, but don't worry,
alpha will be trained along with the rest of your model.
'''
super(Snake,self).__init__()
self.in_features = in_features
# initialize alpha
if alpha is not None:
self.alpha = Parameter(torch.ones(in_features) * alpha) # create a tensor out of alpha
else:
m = Exponential(torch.tensor([0.1]))
self.alpha = Parameter(m.sample(in_features)) # Random init = mix of frequencies
self.alpha.requiresGrad = alpha_trainable # Usually we'll want to train alpha, but maybe for some experiments we won't?
def __init__(self,
input_channel=1,
output_channel=1,
kernel=3,
stride=2,
step=16,
leak=32,
bias=0.5,
winners=0.5,
fodep=None,
w_init=None,
theta=None,
dense=None):
super(ConvColumn, self).__init__()
# model skeleton parameters
self.input_channel = input_channel
self.output_channel = output_channel
self.kernel = kernel
self.stride = stride
# threshold parameters
assert theta or dense, 'either theta or dense should be specified'
self.theta = theta = theta or dense * (kernel * kernel * input_channel)
self.dense = dense = dense or theta / (kernel * kernel * input_channel)
assert dense < 2 * input_channel * kernel * \
kernel, 'invalid theta or density, try setting a smaller value'
w_init = w_init or dense
# spiking control parameters
self.response_function = StepFireLeak(step, leak)
self.fodep = fodep = fodep or self.response_function.fodep
assert fodep >= self.response_function.fodep, f'forced depression should be at least {self.response_function.fodep}'
self.winners = winners
# initialize weight and bias
self.bias = nn.parameter.Parameter(torch.zeros(self.output_channel) +
bias,
requires_grad=False)
self.weight = nn.parameter.Parameter(Exponential(1 / w_init).sample(
(self.output_channel, self.input_channel, self.kernel,
self.kernel)).clip(0, 1),
requires_grad=True)
print(
'Building convolutional connected TNN layer with '
f'theta={theta:.4f}, '
f'dense={dense:.4f}, '
f'fodep={fodep}, ', f'winners={winners}, '
f'bias={bias}')
class Beamforming(EnvWrapper):
def __init__(self, config):
self.config = config
torch.manual_seed(config["random_seed"])
self.M = int(config["M"])
self.K = int(config["K"])
self.G = Exponential(1).sample((self.M, self.K)).float().cpu().detach().numpy()
self.G = self.G[:, np.argsort(self.G.sum(axis=0))]
self.P = float(config["transmission_power"])
def sinr(self, W):
W2 = np.square(W)
gamma = np.zeros(self.K, dtype='float32')
for k in range(self.K):
nom = np.dot(self.G[:,k].reshape(1,self.M), W2[:,k].reshape(self.M,1))
# denom = (self.G.sum(axis=1) - self.G[:,k]).reshape(self.M, 1)
# denom += self.G[:, list(range(k))].sum(axis=1).reshape(self.M, 1)
# denom += (self.G.sum(axis=1) * (self.K-2) + self.G[:,k]).reshape(self.M, 1)
denom = self.G[:, list(range(k))].sum(axis=1)
denom = denom.reshape(self.M, 1)
denom = np.dot(W2[:,k].reshape(1,self.M), denom)
denom += np.float32(1.0/self.P)
gamma[k] = nom/denom
return torch.from_numpy(gamma).float().to("cpu")
def reset(self):
W = torch.rand(self.M, self.K)
W = torch.nn.functional.normalize(W, p=2)
return self.sinr(W)
def step(self, action):
state = self.sinr(action.reshape(self.M, self.K))
reward = (1+state).log2().sum()
return state, reward, 0
def __init__(self,
actor_critic,
clip_param,
ppo_epoch,
num_mini_batch,
value_loss_coef,
entropy_coef,
lr=None,
eps=None,
max_grad_norm=None,
use_clipped_value_loss=True,
ftrl_mode=False,
correlated_mode=False):
self.correlated_mode = correlated_mode # we inject embedding in this case
self.ftrl_mode = ftrl_mode # we use regularization for FTRL.
self.actor_critic = actor_critic
self.clip_param = clip_param
self.ppo_epoch = ppo_epoch
self.num_mini_batch = num_mini_batch
self.value_loss_coef = value_loss_coef
self.entropy_coef = entropy_coef
self.max_grad_norm = max_grad_norm
self.use_clipped_value_loss = use_clipped_value_loss
self.optimizer = optim.Adam(actor_critic.parameters(), lr=lr, eps=eps)
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
params = torch.Tensor().to(self.device)
for param in self.actor_critic.parameters():
params = torch.cat((params, param.view(-1)))
self.randomness = Exponential(torch.ones(len(params))).sample().to(
self.device)
class EmitterSamplerBlinking(EmitterSamplerFrameIndependent):
def __init__(self, *, structure: structure_prior.StructurePrior, intensity_mu_sig: tuple, lifetime: float,
frame_range: tuple, xy_unit: str, px_size: tuple, density=None, em_avg=None, intensity_th=None):
"""
Args:
structure:
intensity_mu_sig:
lifetime:
xy_unit:
px_size:
frame_range: specifies the frame range
density:
em_avg:
intensity_th:
"""
super().__init__(structure=structure,
photon_range=None,
xy_unit=xy_unit,
px_size=px_size,
density=density,
em_avg=em_avg)
self.n_sampler = np.random.poisson
self.frame_range = frame_range
self.intensity_mu_sig = intensity_mu_sig
self.intensity_dist = torch.distributions.normal.Normal(self.intensity_mu_sig[0],
self.intensity_mu_sig[1])
self.intensity_th = intensity_th if intensity_th is not None else 1e-8
self.lifetime_avg = lifetime
self.lifetime_dist = Exponential(1 / self.lifetime_avg) # parse the rate not the scale ...
self.t0_dist = torch.distributions.uniform.Uniform(*self._frame_range_plus)
"""
Determine the total number of emitters. Depends on lifetime, frames and emitters.
(lifetime + 1) because of binning effect.
"""
self._emitter_av_total = self._em_avg * self._num_frames_plus / (self.lifetime_avg + 1)
@property
def _frame_range_plus(self):
"""
Frame range including buffer in front and end to account for build up effects.
"""
return self.frame_range[0] - 3 * self.lifetime_avg, self.frame_range[1] + 3 * self.lifetime_avg
@property
def num_frames(self):
return self.frame_range[1] - self.frame_range[0] + 1
@property
def _num_frames_plus(self):
return self._frame_range_plus[1] - self._frame_range_plus[0] + 1
def sample(self):
"""
Return sampled EmitterSet in the specified frame range.
Returns:
EmitterSet
"""
n = self.n_sampler(self._emitter_av_total)
loose_em = self.sample_loose_emitter(n=n)
em = loose_em.return_emitterset()
em = em.get_subset_frame(*self.frame_range) # because the simulated frame range is larger
return em
def sample_n(self, *args, **kwargs):
raise NotImplementedError
def sample_loose_emitter(self, n) -> decode.generic.emitter.LooseEmitterSet:
"""
Generate loose EmitterSet. Loose emitters are emitters that are not yet binned to frames.
Args:
n: number of 'loose' emitters
Returns:
LooseEmitterSet
"""
xyz = self.structure.sample(n)
"""Draw from intensity distribution but clamp the value so as not to fall below 0."""
intensity = torch.clamp(self.intensity_dist.sample((n,)), self.intensity_th)
"""Distribute emitters in time. Increase the range a bit."""
t0 = self.t0_dist.sample((n,))
ontime = self.lifetime_dist.rsample((n,))
return decode.generic.emitter.LooseEmitterSet(xyz, intensity, ontime, t0, id=torch.arange(n).long(),
xy_unit=self.xy_unit, px_size=self.px_size)
@classmethod
def parse(cls, param, structure, frames: tuple):
return cls(structure=structure,
intensity_mu_sig=param.Simulation.intensity_mu_sig,
lifetime=param.Simulation.lifetime_avg,
xy_unit=param.Simulation.xy_unit,
px_size=param.Camera.px_size,
frame_range=frames,
density=param.Simulation.density,
em_avg=param.Simulation.emitter_av,
intensity_th=param.Simulation.intensity_th)
for param in x_player.parameters(): # the same number of params for y_player
params = torch.cat((params, param.view(-1)))
for t in range(NUM_STEPS):
# player x's turn
x_optimizer.zero_grad()
x_loss = torch.zeros(1, 1)
y = y_player()
for param in x_queue:
x_player.load_state_dict(param)
x = x_player()
x_loss += torch.mm(torch.mm(x, A), y.T)
x_loss /= len(x_queue)
# perturbation
randomness = Exponential(x_rate * torch.ones(len(params))).sample()
params = torch.Tensor()
for param in x_player.parameters():
params = torch.cat((params, param.view(-1)))
x_loss -= torch.dot(params, randomness)
x_loss.backward()
x_optimizer.step()
x_queue, m = queue_update(queue=x_queue,
m=m,
K=K,
t=t + 1,
ft=copy.deepcopy(x_player.state_dict()),
inc=inc)
# player y's turn
def _update_weight(self):
if self.epoch > self.epoch_th:
self.alpha = Exponential(torch.ones([1, self.n_a])).sample()
class NbsRunner(CnnRunner):
def __init__(self, loader, model, optim, lr_scheduler, num_epoch,
loss_with_weight, val_metric, test_metric, logger, model_path,
rank, epoch_th, num_mc):
self.num_mc = num_mc
self.n_a = loader.n_a
self.epoch_th = epoch_th
self.group_indices = loader.groups
self.alpha = torch.ones([1, self.n_a])
super().__init__(loader, model, optim, lr_scheduler, num_epoch,
loss_with_weight, val_metric, test_metric, logger,
model_path, rank)
self.save_kwargs['alpha'] = self.alpha
def _update_weight(self):
if self.epoch > self.epoch_th:
self.alpha = Exponential(torch.ones([1, self.n_a])).sample()
def _calc_loss(self, img, label, idx):
n0 = img.size(0)
u_is = []
for i in idx:
u_i = np.where(self.group_indices == i.item())[0][0]
u_is += [u_i]
w = self.alpha[0, u_is].cuda()
output = self.model(img.cuda(non_blocking=True),
self.alpha.repeat_interleave(n0, 0))
label = label.cuda(non_blocking=True)
loss_ = 0
for loss, w in self.loss_with_weight:
_loss = w * loss(output, label, w)
loss_ += _loss
return loss_
@torch.no_grad()
def _valid_a_batch(self, img, label, _, with_output=False):
self._update_weight()
self.model.eval()
output = self.model(img.cuda(non_blocking=True), self.num_mc)
label = label.cuda(non_blocking=True)
result = self.val_metric(output.mean(0), label)
if with_output:
result = [result, output]
return result
def test(self):
self.load('model.pth')
loader = self.loader.load('test')
if self.rank == 0:
t_iter = tqdm(loader, total=self.loader.len)
else:
t_iter = loader
outputs = []
labels = []
self.model.eval()
for img, label, index in t_iter:
_, output = self._valid_a_batch(img, label, with_output=True)
outputs += [gather_tensor(output).cpu().numpy()]
labels += [gather_tensor(label).cpu().numpy()]
labels = np.concatenate(labels)
outputs = np.concatenate(outputs, axis=1)
acc = (outputs.mean(0).argmax(-1) == labels).mean() * 100
ece = calc_ece(outputs.mean(0), labels)
nll, brier = calc_nll_brier(
outputs.mean(0), labels,
one_hot(torch.from_numpy(labels.astype(int)),
self.model.module.classifer.out_features).numpy())
log = f"[Test] ACC: {acc:.2f}, ECE : {ece:.2f}, "
log += f"NLL : {nll:.2f}, Brier : {brier:.2f}"
self.log(log, 'info')
with h5py.File(f"{self.model_path}/output.h5", 'w') as h:
h.create_dataset('output', data=outputs)
h.create_dataset('label', data=labels)
class NbsRunner(CnnRunner):
def __init__(self, loader, model, optim, lr_scheduler, num_epoch,
loss_with_weight, val_metric, test_metric, logger, model_path,
rank, epoch_th, num_mc, adv_training):
self.num_mc = num_mc
self.n_a = loader.n_a
self.epoch_th = epoch_th
self.alpha = torch.ones([1, self.n_a])
super().__init__(loader, model, optim, lr_scheduler, num_epoch,
loss_with_weight, val_metric, test_metric, logger,
model_path, rank, adv_training)
self.save_kwargs['alpha'] = self.alpha
self._update_weight()
def _update_weight(self):
if self.epoch > self.epoch_th:
self.alpha = Exponential(torch.ones([1, self.n_a])).sample()
def _calc_loss(self, img, label, idx):
n0 = img.size(0)
w = self.alpha[0, idx].cuda()
output = self.model(img.cuda(non_blocking=True),
self.alpha.repeat_interleave(n0, 0))
for _ in range(output.dim() - w.dim()):
w.unsqueeze_(-1)
label = label.cuda(non_blocking=True)
loss_ = 0
for loss, _w in self.loss_with_weight:
_loss = _w * loss(output, label, w)
loss_ += _loss
return loss_
@torch.no_grad()
def _valid_a_batch(self, img, label, with_output=False):
self._update_weight()
self.model.eval()
output = self.model(img.cuda(non_blocking=True), self.num_mc)
label = label.cuda(non_blocking=True)
result = self.val_metric(output.mean(0), label)
if with_output:
result = [result, output]
return result
def test(self, is_seg):
self.load('model.pth')
loader = self.loader.load('test')
if self.rank == 0:
t_iter = tqdm(loader, total=self.loader.len)
else:
t_iter = loader
outputs = []
labels = []
metrics = []
self.model.eval()
for img, label in t_iter:
_metric, output = self._valid_a_batch(img, label, with_output=True)
labels += [gather_tensor(label).cpu().numpy()]
outputs += [gather_tensor(output).cpu().numpy()]
metrics += [gather_tensor(_metric).cpu().numpy()]
if is_seg:
met = np.concatenate(metrics).mean()
self.log(f"[Test] MeanIOU: {met:.2f}", 'info')
save_path = Path(self.model_path) / 'infer'
save_path.mkdir(parents=True, exist_ok=True)
index = 0
for out, label in zip(outputs, labels):
for i in range(label.shape[0]):
l = label[i]
o = out[:, i]
with h5py.File(f"{save_path}/{index}.h5", 'w') as h:
h.create_dataset('output', data=o)
h.create_dataset('label', data=l)
index += 1
else:
labels = np.concatenate(labels)
outputs = np.concatenate(outputs, axis=1)
acc = (outputs.mean(0).argmax(-1) == labels).mean() * 100
ece = calc_ece(softmax(outputs, -1).mean(0), labels)
nll, brier = calc_nll_brier_mc(outputs, labels)
log = f"[Test] ACC: {acc:.2f}, ECE : {ece:.2f}, "
log += f"NLL : {nll:.2f}, Brier : {brier:.2f}"
self.log(log, 'info')
with h5py.File(f"{self.model_path}/output.h5", 'w') as h:
h.create_dataset('output', data=outputs)
h.create_dataset('label', data=labels)
def __init__(self, rate=1):
self.rate = rate
self.exp = Exponential(rate)
