|
import math |
|
from argparse import Namespace |
|
|
|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
|
|
from models import register |
|
|
|
|
|
class gen_basis(nn.Module): |
|
def __init__(self, args): |
|
super(gen_basis, self).__init__() |
|
self.basis_num = args.basis_num |
|
self.hidden = args.hidden |
|
self.state = args.state |
|
self.path=args.path |
|
|
|
def init_basis_bias(self): |
|
self.w0 = nn.Parameter(torch.Tensor(self.basis_num,self.hidden*580), requires_grad=True) |
|
nn.init.kaiming_uniform_(self.w0, a=math.sqrt(5)) |
|
self.w1 = nn.Parameter(torch.Tensor(self.basis_num,self.hidden*self.hidden), requires_grad=True) |
|
nn.init.kaiming_uniform_(self.w1, a=math.sqrt(5)) |
|
self.w2 = nn.Parameter(torch.Tensor(self.basis_num,self.hidden*self.hidden), requires_grad=True) |
|
nn.init.kaiming_uniform_(self.w2, a=math.sqrt(5)) |
|
self.w3 = nn.Parameter(torch.Tensor(self.basis_num,self.hidden*self.hidden), requires_grad=True) |
|
nn.init.kaiming_uniform_(self.w3, a=math.sqrt(5)) |
|
self.w4 = nn.Parameter(torch.Tensor(self.basis_num,3*self.hidden), requires_grad=True) |
|
nn.init.kaiming_uniform_(self.w4, a=math.sqrt(5)) |
|
basis = [self.w0, self.w1, self.w2, self.w3, self.w4] |
|
self.bias1 = nn.Parameter(torch.Tensor(self.basis_num,self.hidden), requires_grad=True) |
|
self.bias2 = nn.Parameter(torch.Tensor(self.basis_num,self.hidden), requires_grad=True) |
|
self.bias3 = nn.Parameter(torch.Tensor(self.basis_num,self.hidden), requires_grad=True) |
|
self.bias4 = nn.Parameter(torch.Tensor(self.basis_num,self.hidden), requires_grad=True) |
|
self.bias5 = nn.Parameter(torch.Tensor(self.basis_num,3), requires_grad=True) |
|
bias = [self.bias1,self.bias2,self.bias3,self.bias4,self.bias5] |
|
|
|
for i in range(len(bias)): |
|
fan_in, _ = nn.init._calculate_fan_in_and_fan_out(basis[i]) |
|
bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 |
|
nn.init.uniform_(bias[i], -bound, bound) |
|
|
|
|
|
|
|
return basis,bias |
|
|
|
|
|
def load_basis_for_test_kaiming(self,path): |
|
model_spec = torch.load(path)['model'] |
|
w0 = model_spec['sd']['basis.w0'] |
|
w1 = model_spec['sd']['basis.w1'] |
|
w2 = model_spec['sd']['basis.w2'] |
|
w3 = model_spec['sd']['basis.w3'] |
|
w4 = model_spec['sd']['basis.w4'] |
|
b0 = model_spec['sd']['basis.bias1'] |
|
b1 = model_spec['sd']['basis.bias2'] |
|
b2 = model_spec['sd']['basis.bias3'] |
|
b3 = model_spec['sd']['basis.bias4'] |
|
b4 = model_spec['sd']['basis.bias5'] |
|
torch.cuda.empty_cache() |
|
return [w0,w1,w2,w3,w4],[b0,b1,b2,b3,b4] |
|
|
|
def forward(self): |
|
if self.state=='train': |
|
print('init_basis_use_kaiming') |
|
res=self.init_basis_bias() |
|
else: |
|
print('load_basis_from_model') |
|
res=self.load_basis_for_test_kaiming(self.path) |
|
return res |
|
|
|
@register('basis') |
|
def make_basis(basis_num=10,hidden=16,state=None,path=None): |
|
args = Namespace() |
|
args.basis_num = basis_num |
|
args.hidden = hidden |
|
args.state = state |
|
args.path = path |
|
return gen_basis(args) |
|
|
