You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
109 lines
3.3 KiB
109 lines
3.3 KiB
import torch
|
|
import pdb
|
|
import torch.nn as nn
|
|
import math
|
|
from torch.autograd import Variable
|
|
from torch.autograd import Function
|
|
|
|
import numpy as np
|
|
|
|
|
|
def Binarize(tensor,quant_mode='det'):
|
|
if quant_mode=='det':
|
|
return tensor.sign()
|
|
else:
|
|
return tensor.add_(1).div_(2).add_(torch.rand(tensor.size()).add(-0.5)).clamp_(0,1).round().mul_(2).add_(-1)
|
|
|
|
|
|
|
|
|
|
class HingeLoss(nn.Module):
|
|
def __init__(self):
|
|
super(HingeLoss,self).__init__()
|
|
self.margin=1.0
|
|
|
|
def hinge_loss(self,input,target):
|
|
#import pdb; pdb.set_trace()
|
|
output=self.margin-input.mul(target)
|
|
output[output.le(0)]=0
|
|
return output.mean()
|
|
|
|
def forward(self, input, target):
|
|
return self.hinge_loss(input,target)
|
|
|
|
class SqrtHingeLossFunction(Function):
|
|
def __init__(self):
|
|
super(SqrtHingeLossFunction,self).__init__()
|
|
self.margin=1.0
|
|
|
|
def forward(self, input, target):
|
|
output=self.margin-input.mul(target)
|
|
output[output.le(0)]=0
|
|
self.save_for_backward(input, target)
|
|
loss=output.mul(output).sum(0).sum(1).div(target.numel())
|
|
return loss
|
|
|
|
def backward(self,grad_output):
|
|
input, target = self.saved_tensors
|
|
output=self.margin-input.mul(target)
|
|
output[output.le(0)]=0
|
|
import pdb; pdb.set_trace()
|
|
grad_output.resize_as_(input).copy_(target).mul_(-2).mul_(output)
|
|
grad_output.mul_(output.ne(0).float())
|
|
grad_output.div_(input.numel())
|
|
return grad_output,grad_output
|
|
|
|
def Quantize(tensor,quant_mode='det', params=None, numBits=8):
|
|
tensor.clamp_(-2**(numBits-1),2**(numBits-1))
|
|
if quant_mode=='det':
|
|
tensor=tensor.mul(2**(numBits-1)).round().div(2**(numBits-1))
|
|
else:
|
|
tensor=tensor.mul(2**(numBits-1)).round().add(torch.rand(tensor.size()).add(-0.5)).div(2**(numBits-1))
|
|
quant_fixed(tensor, params)
|
|
return tensor
|
|
|
|
#import torch.nn._functions as tnnf
|
|
|
|
|
|
class BinarizeLinear(nn.Linear):
|
|
|
|
def __init__(self, *kargs, **kwargs):
|
|
super(BinarizeLinear, self).__init__(*kargs, **kwargs)
|
|
|
|
def forward(self, input):
|
|
|
|
# if input.size(1) != 784:
|
|
# input.data=Binarize(input.data)
|
|
if not hasattr(self.weight,'org'):
|
|
self.weight.org=self.weight.data.clone()
|
|
self.weight.data=Binarize(self.weight.org)
|
|
out = nn.functional.linear(input, self.weight)
|
|
if not self.bias is None:
|
|
self.bias.org=self.bias.data.clone()
|
|
out += self.bias.view(1, -1).expand_as(out)
|
|
|
|
return out
|
|
|
|
class BinarizeConv2d(nn.Conv2d):
|
|
|
|
def __init__(self, *kargs, **kwargs):
|
|
super(BinarizeConv2d, self).__init__(*kargs, **kwargs)
|
|
|
|
|
|
def forward(self, input):
|
|
# if input.size(1) != 3:
|
|
# input.data = Binarize(input.data)
|
|
if not hasattr(self.weight,'org'):
|
|
self.weight.org=self.weight.data.clone()
|
|
self.weight.data=Binarize(self.weight.org)
|
|
|
|
out = nn.functional.conv2d(input, self.weight, None, self.stride,
|
|
self.padding, self.dilation, self.groups)
|
|
|
|
if not self.bias is None:
|
|
self.bias.org=self.bias.data.clone()
|
|
out += self.bias.view(1, -1, 1, 1).expand_as(out)
|
|
|
|
return out
|
|
# x = torch.tensor([[255.0, 200.0, 201.0], [210.0, 222.0, 223.0]])
|
|
# print(Quantize(x,quant_mode='det', params=None, numBits=8))
|