-
Notifications
You must be signed in to change notification settings - Fork 4
/
fixation_duration.py
50 lines (42 loc) · 1.64 KB
/
fixation_duration.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Normal
from utils import WeightedGlobalMeanPool
class FixationDuration(nn.Module):
def __init__(self, feat_dim, roi_hidden_dim):
super(FixationDuration, self).__init__()
self.feat_dim = feat_dim
self.roi_hidden_dim = roi_hidden_dim
self.mu = nn.Sequential(
nn.Linear(self.feat_dim + self.roi_hidden_dim, 128),
nn.Tanh(),
nn.Linear(128, 32),
nn.Tanh(),
nn.Linear(32, 1)
)
self.mu[-1].bias.data.copy_(torch.rand_like(self.mu[-1].bias))
self.log_var = nn.Sequential(
nn.Linear(self.feat_dim + self.roi_hidden_dim, 128),
nn.Tanh(),
nn.Linear(128, 32),
nn.Tanh(),
nn.Linear(32, 1)
)
self.pool = WeightedGlobalMeanPool(keepdim=False)
def forward(self, feat, hidden, current_ROI):
spatial_size = feat.size()[2:]
with torch.no_grad():
current_roi = F.interpolate(current_ROI, size=spatial_size, mode='bilinear')
feat_vec = self.pool(feat, current_roi)
hidden_vec = self.pool(hidden, current_roi)
gaussian_input = torch.cat([feat_vec, hidden_vec], dim=-1)
mu = self.mu(gaussian_input)
mu = torch.clamp(mu, 0, 4)
sigma_squared = torch.exp(self.log_var(gaussian_input))
sigma = torch.sqrt(sigma_squared)
sigma = sigma.clamp(0.05, 0.1)
normal_distribution = Normal(mu, sigma)
sample = normal_distribution.rsample()
sample = sample.clamp(0, 4)
return sample