mirror of https://github.com/menyifang/DCT-Net
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
15 lines
379 B
Python
15 lines
379 B
Python
2 years ago
|
import torch
|
||
|
|
from torch import nn
|
||
|
|
|
||
|
|
|
||
|
|
class WNormLoss(nn.Module):
|
||
|
|
|
||
|
|
def __init__(self, start_from_latent_avg=True):
|
||
|
|
super(WNormLoss, self).__init__()
|
||
|
|
self.start_from_latent_avg = start_from_latent_avg
|
||
|
|
|
||
|
|
def forward(self, latent, latent_avg=None):
|
||
|
|
if self.start_from_latent_avg:
|
||
|
|
latent = latent - latent_avg
|
||
|
|
return torch.sum(latent.norm(2, dim=(1, 2))) / latent.shape[0]
|