Archives/imaginAIry
2
0
Fork 0
mirror of https://github.com/brycedrennan/imaginAIry synced 2024-11-17 09:25:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

lint

Browse Source
This commit is contained in:
Bryce 2022-09-21 09:14:39 -07:00
parent ff7fc0dab7
commit e23e363bf5
1 changed files with 17 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
imaginairy/modules/distributions.py
View File

@ -24,23 +24,24 @@ class DiagonalGaussianDistribution:
def kl(self, other=None):
if self.deterministic:
return torch.Tensor([0.0])
else:
if other is None:
return 0.5 * torch.sum(
torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
dim=[1, 2, 3],
)
else:
return 0.5 * torch.sum(
torch.pow(self.mean - other.mean, 2) / other.var
+ self.var / other.var
- 1.0
- self.logvar
+ other.logvar,
dim=[1, 2, 3],
)
def nll(self, sample, dims=[1, 2, 3]):
if other is None:
return 0.5 * torch.sum(
torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
dim=[1, 2, 3],
)
return 0.5 * torch.sum(
torch.pow(self.mean - other.mean, 2) / other.var
+ self.var / other.var
- 1.0
- self.logvar
+ other.logvar,
dim=[1, 2, 3],
)
def nll(self, sample, dims=None):
dims = dims if dims is None else [1, 2, 3]
if self.deterministic:
return torch.Tensor([0.0])
logtwopi = np.log(2.0 * np.pi)