I have made a partially connected layer, in which input layer and output layer are not completely connected according to adjacency matrix.
ex.
input layer [in_features]
weight matrix [out_features,in_features] : only specific element is non-zero and learnable.
bias [out_features]
->
output layer [out_features]
However, It has a memory leakage problem.
How can I index and update specific elements in tensor without making a new tensor.
Below is my code.
from torch.nn.parameter import Parameter
import torch.nn.init as init
import torch.nn.functional as F
import torch.nn as nn
import math
import numpy as np
import torch
def kaiming_uniform_fan_(tensor, fan=None, a=0, nonlinearity='leaky_relu'):
gain = init.calculate_gain(nonlinearity, a)
std = gain / math.sqrt(fan)
bound = math.sqrt(3.0) * std # Calculate uniform bounds from standard deviation
with torch.no_grad():
return tensor.uniform_(-bound, bound)
class Linear_Partial(nn.Module):
__constants__ = ['bias', 'adjacency']
def __init__(self, adjacency,bias=True):
super(Linear_Partial, self).__init__()
# adjacency matrix: non-zero indice are only learnable
assert ((adjacency!=1) & (adjacency!=0)).sum()==0, 'value not in [0,1]'
self.adjacency = Parameter(torch.Tensor(adjacency))
self.in_features=adjacency.shape[1]
self.out_features=adjacency.shape[0]
# elements in weights which are learnable.
self.weight_learnable_count=(adjacency!=0).sum()
self.weight_learnable=Parameter(torch.Tensor(self.weight_learnable_count))
self.weight=torch.zeros([self.out_features,self.in_features])
self.update_weight_manually()
#self.weight[adjacency==1] = self.weight_learnable
#self.register_parameter('weight1',self.weight)
if bias:
self.bias = Parameter(torch.Tensor(self.out_features))
else:
self.register_parameter('bias', None)
self.reset_parameters()
# since weight is not registered as parameter, we have to update weight value manually
def update_weight_manually(self):
self.weight[self.adjacency==1] = self.weight_learnable
def reset_parameters(self):
kaiming_uniform_fan_(self.weight_learnable, fan=self.in_features, a=math.sqrt(5))
self.update_weight_manually()
#self.weight[self.adjacency==0] = 0
if self.bias is not None:
fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
bound = 1 / math.sqrt(fan_in)
init.uniform_(self.bias, -bound, bound)
def forward(self, input):
self.update_weight_manually()
return F.linear(input, self.weight, self.bias)
def extra_repr(self):
return 'in_features={}, out_features={}, bias={}'.format(
self.in_features, self.out_features, self.bias is not None
)
adjacency=(np.random.rand(20000,10000)>0.5).astype(int)
model=Linear_Partial(adjacency)
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
model.to(device)
input_sample=torch.ones(10000).to(device)
def init_weights(m,device):
if type(m) == Linear_Partial:
m.weight=m.weight.to(device)
print("Sent to {}: {}".format(device,m))
if True:
model.apply(lambda m: init_weights(m,device))
output_sample=model(input_sample)
Error message
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
<ipython-input-3-4099065c43f2> in <module>
----> 1 output_sample=model(input_sample)
~/tools/miniconda3/envs/pytorch/lib/python3.6/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
539 result = self._slow_forward(*input, **kwargs)
540 else:
--> 541 result = self.forward(*input, **kwargs)
542 for hook in self._forward_hooks.values():
543 hook_result = hook(self, input, result)
<ipython-input-1-fb1702069a84> in forward(self, input)
57
58 def forward(self, input):
---> 59 self.update_weight_manually()
60 return F.linear(input, self.weight, self.bias)
61
<ipython-input-1-fb1702069a84> in update_weight_manually(self)
44
45 def update_weight_manually(self):
---> 46 self.weight[self.adjacency==1] = self.weight_learnable
47
48 def reset_parameters(self):
RuntimeError: CUDA out of memory. Tried to allocate 2.98 GiB (GPU 1; 7.77 GiB total capacity; 2.24 GiB already allocated; 2.02 GiB free; 2.80 GiB cached)