File size: 2,131 Bytes
96170c3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
import copy
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import Parameter
class LinearSVDO(nn.Module):
def __init__(self, in_features, out_features, alpha_threshold, theta_threshold, device):
super(LinearSVDO, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.alpha_threshold = alpha_threshold
self.theta_threshold = theta_threshold
self.device = device
self.W = Parameter(torch.Tensor(out_features, in_features))
self.log_sigma = Parameter(torch.Tensor(out_features, in_features))
self.bias = Parameter(torch.Tensor(1, out_features))
self.reset_parameters()
def reset_parameters(self):
self.bias.data.zero_()
self.W.data.normal_(0, 0.02)
self.log_sigma.data.fill_(-5)
def forward(self, x):
self.log_alpha = self.log_sigma * 2.0 - 2.0 * torch.log(1e-16 + torch.abs(self.W))
self.log_alpha = torch.clamp(self.log_alpha, -10, 10)
if self.training:
lrt_mean = F.linear(x, self.W) + self.bias
lrt_std = F.linear(torch.sqrt(x * x), torch.exp(2*self.log_sigma)+ 1e-8)
eps = torch.randn_like(lrt_std)
return lrt_mean + lrt_std * eps
out = self.W * (self.log_alpha < self.alpha_threshold).float()
out = F.linear(x, out) + self.bias
return out
def get_pruned_weights(self):
W = self.W * (self.log_alpha < self.alpha_threshold).float()
return W
def get_num_remained_weights(self):
num = ((self.log_alpha < self.alpha_threshold) * (torch.abs(self.W) > self.theta_threshold)).sum().item()
return num
def kl_reg(self):
k1, k2, k3 = torch.FloatTensor([0.63576]).to(self.device), torch.FloatTensor([1.8732]).to(self.device), torch.FloatTensor([1.48695]).to(self.device)
KL = k1 * torch.sigmoid(k2 + k3 * self.log_alpha) - 0.5 * torch.log1p(torch.exp(-self.log_alpha))
KL = - torch.sum(KL)
return KL
|