"vscode:/vscode.git/clone" did not exist on "27ecc17a1ea93b5b6b68145df06094de0aa53356"
__init__.py 2.38 KB
Newer Older
Hang Zhang's avatar
init  
Hang Zhang committed
1
2
3
4
5
6
7
8
9
10
11
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created by: Hang Zhang
## ECE Department, Rutgers University
## Email: zhang.hang@rutgers.edu
## Copyright (c) 2017
##
## This source code is licensed under the MIT-style license found in the
## LICENSE file in the root directory of this source tree 
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

import torch
Hang Zhang's avatar
Hang Zhang committed
12
import torch.nn as nn
Hang Zhang's avatar
tested  
Hang Zhang committed
13
from torch.autograd import Function
Hang Zhang's avatar
init  
Hang Zhang committed
14
15
16
17
18
from ._ext import encoding_lib

class aggregate(Function):
	def forward(self, A, R):
		# A \in(BxNxK) R \in(BxNxKxD) => E \in(BxNxD)
Hang Zhang's avatar
Hang Zhang committed
19
		self.save_for_backward(A, R)
Hang Zhang's avatar
init  
Hang Zhang committed
20
21
22
23
24
25
		B, N, K, D = R.size()
		E = A.new(B,K,D)
		# TODO support cpu backend
		encoding_lib.Encoding_Float_aggregate_forward(E, A, R)
		return E

Hang Zhang's avatar
Hang Zhang committed
26
27
	def backward(self, gradE):
		A, R = self.saved_tensors
Hang Zhang's avatar
Hang Zhang committed
28
29
		gradA = A.new().resize_as_(A)
		gradR = R.new().resize_as_(R)
Hang Zhang's avatar
Hang Zhang committed
30
31
32
		encoding_lib.Encoding_Float_aggregate_backward(gradA, gradR, gradE, 
						A, R)
		return gradA, gradR
Hang Zhang's avatar
init  
Hang Zhang committed
33
34


Hang Zhang's avatar
Hang Zhang committed
35
class Aggregate(nn.Module):
Hang Zhang's avatar
init  
Hang Zhang committed
36
37
	def forward(self, A, R):
		return aggregate()(A, R)
Hang Zhang's avatar
Hang Zhang committed
38

Hang Zhang's avatar
Hang Zhang committed
39

Hang Zhang's avatar
Hang Zhang committed
40
41
42
43
44
45
46
47
48
49
50
class Encoding(nn.Module):
	def __init__(self, D, K):
		super(Encoding, self).__init__()
		# init codewords and smoothing factor
		self.D, self.K = D, K
		self.codewords = nn.Parameter(torch.Tensor(K, D), requires_grad=True)
		self.scale = nn.Parameter(torch.Tensor(K), requires_grad=True) 
		self.softmax = nn.Softmax()
		self.reset_params()
		
	def reset_params(self):
Hang Zhang's avatar
Hang Zhang committed
51
52
53
54
		std1 = 1./((self.K*self.D)**(1/2))
		std2 = 1./((self.K)**(1/2))
		self.codewords.data.uniform_(-std1, std1)
		self.scale.data.uniform_(-std2, std2)
Hang Zhang's avatar
Hang Zhang committed
55
56
57
58

	def forward(self, X):
		# input X is a 4D tensor
		assert(X.size(1)==self.D,"Encoding Layer incompatible input channels!")
Hang Zhang's avatar
Hang Zhang committed
59
60
61
62
63
		unpacked = False
		if X.dim() == 3:
			unpacked = True
			X = X.unsqueeze(0)

Hang Zhang's avatar
Hang Zhang committed
64
65
66
67
68
69
70
71
72
73
74
75
76
		B, N, K, D = X.size(0), X.size(2)*X.size(3), self.K, self.D
		# reshape input
		X = X.view(B,D,-1).transpose(1,2)
		# calculate residuals
		R = X.contiguous().view(B,N,1,D).expand(B,N,K,D) - self.codewords.view(
					1,1,K,D).expand(B,N,K,D)
		# assignment weights
		A = R
		A = A.pow(2).sum(3).view(B,N,K)
		A = A*self.scale.view(1,1,K).expand_as(A)
		A = self.softmax(A.view(B*N,K)).view(B,N,K)
		# aggregate
		E = aggregate()(A, R)
Hang Zhang's avatar
Hang Zhang committed
77
78
79

		if unpacked:
			E = E.squeeze(0)
Hang Zhang's avatar
Hang Zhang committed
80
81
82
83
84
		return E

	def __repr__(self):
		return self.__class__.__name__ + '(' \
			+ 'N x ' + str(self.D) + '=>' + str(self.K) + 'x' + str(self.D) + ')'