- dist = torch.distributions.categorical.Categorical(logits=logits)
- t_next = dist.sample()
- input[:, s] = ar_mask[:, s] * t_next + (1 - ar_mask[:, s]) * input[:, s]
+ bs = self.trunk(bs)
+ bs = self.readout(bs)
+ return bs
+ else:
+ bs = self.trunk[start_layer:](bs)
+ bs = self.trunk(bs)
+ bs = self.readout(bs)
+ return bs
+
+ def reset_transformations(self):
+ self.temperature = 1.0
+ for m in self.modules():
+ if isinstance(m, NoiseInjector):
+ m.noise_std = 0.0
+
+ def set_noise_injection(self, noise_std, identifier=None):
+ for m in self.modules():
+ if isinstance(m, NoiseInjector):
+ if identifier is None or identifier == m.identifier:
+ m.noise_std = noise_std
+
+ def record_attention(self, v=True):
+ for m in self.modules():
+ if isinstance(m, QKVAttention):
+ m.record_attention = v
+
+ def retrieve_attention(self):
+ a = []
+ for m in self.modules():
+ if isinstance(m, QKVAttention):
+ a.append(m.a)
+ return a