def loadModel(ckpt_file):
loadedcheckpoint = torch.load(
ckpt_file) #,map_location={'cuda:1':'cuda:0'})
stateDict = loadedcheckpoint['state_dict']
# load parameters
Dtheta = stateDict['l1.theta']
Drr = stateDict['l1.rr']
Gamma = stateDict['l1.gamma']
model = SClayer(Drr, Dtheta, Gamma, T)
model.cuda(gpu_id)
Drr = Variable(Drr.cuda(gpu_id))
Dtheta = Variable(Dtheta.cuda(gpu_id))
dictionary = creatRealDictionary(N_FRAME, Drr, Dtheta)
return model, dictionary
def forward(self, x):
dic = creatRealDictionary(self.T + self.PRE, self.rr, self.theta)
result = torch.matmul(dic, x)
return result
def forward(self, x):
dic = creatRealDictionary(self.T, self.rr, self.theta)
sparsecode = fista(dic, x, 0.01, 80)
return Variable(sparsecode)