brainscore-vision 2.1__py3-none-any.whl

brainscore_vision/models/voneresnet_50_non_stochastic/vonenet/vonenet/__init__.py

@@ -0,0 +1,71 @@
+import torch
+import torch.nn as nn
+import os
+import requests
+
+from .vonenet import VOneNet
+from torch.nn import Module
+
+FILE_WEIGHTS = {'alexnet': 'vonealexnet_e70.pth.tar', 'resnet50': 'voneresnet50_e70.pth.tar',
+                'resnet50_at': 'voneresnet50_at_e96.pth.tar', 'cornets': 'vonecornets_e70.pth.tar',
+                'resnet50_ns': 'voneresnet50_ns_e70.pth.tar'}
+
+
+class Wrapper(Module):
+    def __init__(self, model):
+        super(Wrapper, self).__init__()
+        self.module = model
+
+
+def get_model(model_arch='resnet50', pretrained=True, map_location='cpu', **kwargs):
+    """
+    Returns a VOneNet model.
+    Select pretrained=True for returning one of the 3 pretrained models.
+    model_arch: string with identifier to choose the architecture of the back-end (resnet50, cornets, alexnet)
+    """
+    if pretrained and model_arch:
+        url = f'https://vonenet-models.s3.us-east-2.amazonaws.com/{FILE_WEIGHTS[model_arch.lower()]}'
+        home_dir = os.environ['HOME']
+        vonenet_dir = os.path.join(home_dir, '.vonenet')
+        weightsdir_path = os.path.join(vonenet_dir, FILE_WEIGHTS[model_arch.lower()])
+        if not os.path.exists(vonenet_dir):
+            os.makedirs(vonenet_dir)
+        if not os.path.exists(weightsdir_path):
+            print('Downloading model weights to ', weightsdir_path)
+            r = requests.get(url, allow_redirects=True)
+            open(weightsdir_path, 'wb').write(r.content)
+
+        ckpt_data = torch.load(weightsdir_path, map_location=map_location)
+
+        stride = ckpt_data['flags']['stride']
+        simple_channels = ckpt_data['flags']['simple_channels']
+        complex_channels = ckpt_data['flags']['complex_channels']
+        k_exc = ckpt_data['flags']['k_exc']
+
+        noise_mode = ckpt_data['flags']['noise_mode']
+        noise_scale = ckpt_data['flags']['noise_scale']
+        noise_level = ckpt_data['flags']['noise_level']
+
+        model_id = ckpt_data['flags']['arch'].replace('_','').lower()
+
+        model = globals()[f'VOneNet'](model_arch=model_id, stride=stride, k_exc=k_exc,
+                                      simple_channels=simple_channels, complex_channels=complex_channels,
+                                      noise_mode=noise_mode, noise_scale=noise_scale, noise_level=noise_level)
+
+        if model_arch.lower() == 'resnet50_at':
+            ckpt_data['state_dict'].pop('vone_block.div_u.weight')
+            ckpt_data['state_dict'].pop('vone_block.div_t.weight')
+            model.load_state_dict(ckpt_data['state_dict'])
+        else:
+            model = Wrapper(model)
+            model.load_state_dict(ckpt_data['state_dict'])
+            model = model.module
+
+        model = nn.DataParallel(model)
+    else:
+        model = globals()[f'VOneNet'](model_arch=model_arch, **kwargs)
+        model = nn.DataParallel(model)
+
+    model.to(map_location)
+    return model
+

brainscore_vision/models/voneresnet_50_non_stochastic/vonenet/vonenet/back_ends.py

@@ -0,0 +1,337 @@
+
+import numpy as np
+import torch
+from torch import nn
+from collections import OrderedDict
+
+
+# AlexNet Back-End architecture
+# Based on Torchvision implementation in
+# https://github.com/pytorch/vision/blob/master/torchvision/models/alexnet.py
+class AlexNetBackEnd(nn.Module):
+    def __init__(self, num_classes=1000):
+        super().__init__()
+        self.features = nn.Sequential(
+            nn.Conv2d(64, 192, kernel_size=5, stride=2, padding=2),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
+            nn.Conv2d(192, 384, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(384, 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
+        )
+        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
+        self.classifier = nn.Sequential(
+            nn.Dropout(),
+            nn.Linear(256 * 7 * 7, 4096),
+            nn.ReLU(inplace=True),
+            nn.Dropout(),
+            nn.Linear(4096, 4096),
+            nn.ReLU(inplace=True),
+            nn.Linear(4096, num_classes),
+        )
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.classifier(x)
+        return x
+
+
+# ResNet Back-End architecture
+# Based on Torchvision implementation in
+# https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+    __constants__ = ['downsample']
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(BasicBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True) #
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+    __constants__ = ['downsample']
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(Bottleneck, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(planes * (base_width / 64.)) * groups
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(inplanes, width)
+        self.bn1 = norm_layer(width)
+        self.conv2 = conv3x3(width, width, stride, groups, dilation)
+        self.bn2 = norm_layer(width)
+        self.conv3 = conv1x1(width, planes * self.expansion)
+        self.bn3 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True) # inplace=True
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNetBackEnd(nn.Module):
+    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
+                 groups=1, width_per_group=64, replace_stride_with_dilation=None,
+                 norm_layer=None):
+        super(ResNetBackEnd, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+
+        self.inplanes = 64
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            # each element in the tuple indicates if we should replace
+            # the 2x2 stride with a dilated convolution instead
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError("replace_stride_with_dilation should be None "
+                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+        self.groups = groups
+        self.base_width = width_per_group
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
+                                       dilate=replace_stride_with_dilation[0])
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
+                                       dilate=replace_stride_with_dilation[1])
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
+                                       dilate=replace_stride_with_dilation[2])
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
+                            self.base_width, previous_dilation, norm_layer))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes, groups=self.groups,
+                                base_width=self.base_width, dilation=self.dilation,
+                                norm_layer=norm_layer))
+
+        return nn.Sequential(*layers)
+
+    def _forward_impl(self, x):
+        # See note [TorchScript super()]
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+
+        return x
+
+    def forward(self, x):
+        return self._forward_impl(x)
+
+
+# CORnet-S Back-End architecture
+# Based on CORnet code in
+# https://github.com/dicarlolab/CORnet
+class Flatten(nn.Module):
+    def forward(self, x):
+        return x.view(x.size(0), -1)
+
+
+class Identity(nn.Module):
+    def forward(self, x):
+        return x
+
+
+class CORblock_S(nn.Module):
+
+    scale = 4  # scale of the bottleneck convolution channels
+
+    def __init__(self, in_channels, out_channels, times=1):
+        super().__init__()
+
+        self.times = times
+
+        self.conv_input = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
+        self.skip = nn.Conv2d(out_channels, out_channels,
+                              kernel_size=1, stride=2, bias=False)
+        self.norm_skip = nn.BatchNorm2d(out_channels)
+
+        self.conv1 = nn.Conv2d(out_channels, out_channels * self.scale,
+                               kernel_size=1, bias=False)
+        self.nonlin1 = nn.ReLU(inplace=True) #
+
+        self.conv2 = nn.Conv2d(out_channels * self.scale, out_channels * self.scale,
+                               kernel_size=3, stride=2, padding=1, bias=False)
+        self.nonlin2 = nn.ReLU(inplace=True) #
+
+        self.conv3 = nn.Conv2d(out_channels * self.scale, out_channels,
+                               kernel_size=1, bias=False)
+        self.nonlin3 = nn.ReLU(inplace=True) #
+
+        self.output = Identity()  # for an easy access to this block's output
+
+        # need BatchNorm for each time step for training to work well
+        for t in range(self.times):
+            setattr(self, f'norm1_{t}', nn.BatchNorm2d(out_channels * self.scale))
+            setattr(self, f'norm2_{t}', nn.BatchNorm2d(out_channels * self.scale))
+            setattr(self, f'norm3_{t}', nn.BatchNorm2d(out_channels))
+
+    def forward(self, inp):
+        x = self.conv_input(inp)
+
+        for t in range(self.times):
+            if t == 0:
+                skip = self.norm_skip(self.skip(x))
+                self.conv2.stride = (2, 2)
+            else:
+                skip = x
+                self.conv2.stride = (1, 1)
+
+            x = self.conv1(x)
+            x = getattr(self, f'norm1_{t}')(x)
+            x = self.nonlin1(x)
+
+            x = self.conv2(x)
+            x = getattr(self, f'norm2_{t}')(x)
+            x = self.nonlin2(x)
+
+            x = self.conv3(x)
+            x = getattr(self, f'norm3_{t}')(x)
+
+            x += skip
+            x = self.nonlin3(x)
+            output = self.output(x)
+
+        return output
+
+
+class CORnetSBackEnd(nn.Module):
+    def __init__(self, num_classes=1000):
+        super(CORnetSBackEnd, self).__init__()
+
+        self.V2 = CORblock_S(64, 128, times=2)
+        self.V4 = CORblock_S(128, 256, times=4)
+        self.IT = CORblock_S(256, 512, times=2)
+        self.decoder = nn.Sequential(OrderedDict([
+            ('avgpool', nn.AdaptiveAvgPool2d(1)),
+            ('flatten', Flatten()),
+            ('linear', nn.Linear(512, num_classes)),
+            ('output', Identity())
+        ]))
+
+        # weight initialization
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, np.sqrt(2. / n))
+            # nn.Linear is missing here because I originally forgot
+            # to add it during the training of this network
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def forward(self, x):
+        x = self.V2(x)
+        x = self.V4(x)
+        x = self.IT(x)
+        x = self.decoder(x)
+        return x

brainscore_vision/models/voneresnet_50_non_stochastic/vonenet/vonenet/modules.py

@@ -0,0 +1,126 @@
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+from .utils import gabor_kernel
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class Identity(nn.Module):
+    def forward(self, x):
+        return x
+
+
+class GFB(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=4):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = (kernel_size, kernel_size)
+        self.stride = (stride, stride)
+        self.padding = (kernel_size // 2, kernel_size // 2)
+
+        # Param instatiations
+        self.weight = torch.zeros((out_channels, in_channels, kernel_size, kernel_size))
+
+    def forward(self, x):
+        return F.conv2d(x, self.weight, None, self.stride, self.padding)
+
+    def initialize(self, sf, theta, sigx, sigy, phase):
+        random_channel = torch.randint(0, self.in_channels, (self.out_channels,))
+        for i in range(self.out_channels):
+            self.weight[i, random_channel[i]] = gabor_kernel(frequency=sf[i], sigma_x=sigx[i], sigma_y=sigy[i],
+                                                             theta=theta[i], offset=phase[i], ks=self.kernel_size[0])
+        self.weight = nn.Parameter(self.weight, requires_grad=False)
+
+
+class VOneBlock(nn.Module):
+    def __init__(self, sf, theta, sigx, sigy, phase,
+                 k_exc=25, noise_mode=None, noise_scale=1, noise_level=1,
+                 simple_channels=128, complex_channels=128, ksize=25, stride=4, input_size=224):
+        super().__init__()
+
+        self.in_channels = 3
+
+        self.simple_channels = simple_channels
+        self.complex_channels = complex_channels
+        self.out_channels = simple_channels + complex_channels
+        self.stride = stride
+        self.input_size = input_size
+
+        self.sf = sf
+        self.theta = theta
+        self.sigx = sigx
+        self.sigy = sigy
+        self.phase = phase
+        self.k_exc = k_exc
+
+        self.set_noise_mode(noise_mode, noise_scale, noise_level)
+        self.fixed_noise = None
+
+        self.simple_conv_q0 = GFB(self.in_channels, self.out_channels, ksize, stride)
+        self.simple_conv_q1 = GFB(self.in_channels, self.out_channels, ksize, stride)
+        self.simple_conv_q0.initialize(sf=self.sf, theta=self.theta, sigx=self.sigx, sigy=self.sigy,
+                                       phase=self.phase)
+        self.simple_conv_q1.initialize(sf=self.sf, theta=self.theta, sigx=self.sigx, sigy=self.sigy,
+                                       phase=self.phase + np.pi / 2)
+
+        self.simple = nn.ReLU(inplace=True)
+        self.complex = Identity()
+        self.gabors = Identity()
+        self.noise = nn.ReLU(inplace=True)
+        self.output = Identity()
+
+    def forward(self, x):
+        # Gabor activations [Batch, out_channels, H/stride, W/stride]
+        x = self.gabors_f(x)
+        # Noise [Batch, out_channels, H/stride, W/stride]
+        x = self.noise_f(x)
+        # V1 Block output: (Batch, out_channels, H/stride, W/stride)
+        x = self.output(x)
+        return x
+
+    def gabors_f(self, x):
+        s_q0 = self.simple_conv_q0(x)
+        s_q1 = self.simple_conv_q1(x)
+        c = self.complex(torch.sqrt(s_q0[:, self.simple_channels:, :, :] ** 2 +
+                                    s_q1[:, self.simple_channels:, :, :] ** 2) / np.sqrt(2))
+        s = self.simple(s_q0[:, 0:self.simple_channels, :, :])
+        return self.gabors(self.k_exc * torch.cat((s, c), 1))
+
+    def noise_f(self, x):
+        if self.noise_mode == 'neuronal':
+            eps = 10e-5
+            x *= self.noise_scale
+            x += self.noise_level
+            if self.fixed_noise is not None:
+                x += self.fixed_noise * torch.sqrt(F.relu(x.clone()) + eps)
+            else:
+                x += torch.distributions.normal.Normal(torch.zeros_like(x), scale=1).rsample() * \
+                     torch.sqrt(F.relu(x.clone()) + eps)
+            x -= self.noise_level
+            x /= self.noise_scale
+        if self.noise_mode == 'gaussian':
+            if self.fixed_noise is not None:
+                x += self.fixed_noise * self.noise_scale
+            else:
+                x += torch.distributions.normal.Normal(torch.zeros_like(x), scale=1).rsample() * self.noise_scale
+        return self.noise(x)
+
+    def set_noise_mode(self, noise_mode=None, noise_scale=1, noise_level=1):
+        self.noise_mode = noise_mode
+        self.noise_scale = noise_scale
+        self.noise_level = noise_level
+
+    def fix_noise(self, batch_size=256, seed=None):
+        noise_mean = torch.zeros(batch_size, self.out_channels, int(self.input_size/self.stride),
+                                 int(self.input_size/self.stride))
+        if seed:
+            torch.manual_seed(seed)
+        if self.noise_mode:
+            self.fixed_noise = torch.distributions.normal.Normal(noise_mean, scale=1).rsample().to(device)
+
+    def unfix_noise(self):
+        self.fixed_noise = None

brainscore_vision/models/voneresnet_50_non_stochastic/vonenet/vonenet/params.py

@@ -0,0 +1,100 @@
+
+import numpy as np
+from .utils import sample_dist
+import scipy.stats as stats
+
+
+def generate_gabor_param(features, seed=0, rand_flag=False, sf_corr=0, sf_max=9, sf_min=0):
+    # Generates random sample
+    np.random.seed(seed)
+
+    phase_bins = np.array([0, 360])
+    phase_dist = np.array([1])
+
+    if rand_flag:
+        print('Uniform gabor parameters')
+        ori_bins = np.array([0, 180])
+        ori_dist = np.array([1])
+
+        nx_bins = np.array([0.1, 10**0.2])
+        nx_dist = np.array([1])
+
+        ny_bins = np.array([0.1, 10**0.2])
+        ny_dist = np.array([1])
+
+        # sf_bins = np.array([0.5, 8])
+        # sf_dist = np.array([1])
+
+        sf_bins = np.array([0.5, 0.7, 1.0, 1.4, 2.0, 2.8, 4.0, 5.6, 8])
+        sf_dist = np.array([1,  1,  1, 1, 1, 1, 1, 1])
+
+        sfmax_ind = np.where(sf_bins < sf_max)[0][-1]
+        sfmin_ind = np.where(sf_bins >= sf_min)[0][0]
+
+        sf_bins = sf_bins[sfmin_ind:sfmax_ind+1]
+        sf_dist = sf_dist[sfmin_ind:sfmax_ind]
+
+        sf_dist = sf_dist / sf_dist.sum()
+    else:
+        print('Neuronal distributions gabor parameters')
+        # DeValois 1982a
+        ori_bins = np.array([-22.5, 22.5, 67.5, 112.5, 157.5])
+        ori_dist = np.array([66, 49, 77, 54])
+        ori_dist = ori_dist / ori_dist.sum()
+
+        # Schiller 1976
+        cov_mat = np.array([[1, sf_corr], [sf_corr, 1]])
+
+        # Ringach 2002b
+        nx_bins = np.logspace(-1, 0.2, 6, base=10)
+        ny_bins = np.logspace(-1, 0.2, 6, base=10)
+        n_joint_dist = np.array([[2.,  0.,  1.,  0.,  0.],
+                                 [8.,  9.,  4.,  1.,  0.],
+                                 [1.,  2., 19., 17.,  3.],
+                                 [0.,  0.,  1.,  7.,  4.],
+                                 [0.,  0.,  0.,  0.,  0.]])
+        n_joint_dist = n_joint_dist / n_joint_dist.sum()
+        nx_dist = n_joint_dist.sum(axis=1)
+        nx_dist = nx_dist / nx_dist.sum()
+        ny_dist_marg = n_joint_dist / n_joint_dist.sum(axis=1, keepdims=True)
+
+        # DeValois 1982b
+        sf_bins = np.array([0.5, 0.7, 1.0, 1.4, 2.0, 2.8, 4.0, 5.6, 8])
+        sf_dist = np.array([4,  4,  8, 25, 32, 26, 28, 12])
+
+        sfmax_ind = np.where(sf_bins <= sf_max)[0][-1]
+        sfmin_ind = np.where(sf_bins >= sf_min)[0][0]
+
+        sf_bins = sf_bins[sfmin_ind:sfmax_ind+1]
+        sf_dist = sf_dist[sfmin_ind:sfmax_ind]
+
+        sf_dist = sf_dist / sf_dist.sum()
+
+    phase = sample_dist(phase_dist, phase_bins, features)
+    ori = sample_dist(ori_dist, ori_bins, features)
+    ori[ori < 0] = ori[ori < 0] + 180
+
+    if rand_flag:
+        sf = sample_dist(sf_dist, sf_bins, features, scale='log2')
+        nx = sample_dist(nx_dist, nx_bins, features, scale='log10')
+        ny = sample_dist(ny_dist, ny_bins, features, scale='log10')
+    else:
+
+        samps = np.random.multivariate_normal([0, 0], cov_mat, features)
+        samps_cdf = stats.norm.cdf(samps)
+
+        nx = np.interp(samps_cdf[:,0], np.hstack(([0], nx_dist.cumsum())), np.log10(nx_bins))
+        nx = 10**nx
+
+        ny_samp = np.random.rand(features)
+        ny = np.zeros(features)
+        for samp_ind, nx_samp in enumerate(nx):
+            bin_id = np.argwhere(nx_bins < nx_samp)[-1]
+            ny[samp_ind] = np.interp(ny_samp[samp_ind], np.hstack(([0], ny_dist_marg[bin_id, :].cumsum())),
+                                             np.log10(ny_bins))
+        ny = 10**ny
+
+        sf = np.interp(samps_cdf[:,1], np.hstack(([0], sf_dist.cumsum())), np.log2(sf_bins))
+        sf = 2**sf
+
+    return sf, ori, phase, nx, ny

brainscore_vision/models/voneresnet_50_non_stochastic/vonenet/vonenet/utils.py

@@ -0,0 +1,32 @@
+
+import numpy as np
+import torch
+
+
+def gabor_kernel(frequency,  sigma_x, sigma_y, theta=0, offset=0, ks=61):
+
+    w = ks // 2
+    grid_val = torch.arange(-w, w+1, dtype=torch.float)
+    x, y = torch.meshgrid(grid_val, grid_val)
+    rotx = x * np.cos(theta) + y * np.sin(theta)
+    roty = -x * np.sin(theta) + y * np.cos(theta)
+    g = torch.zeros(y.shape)
+    g[:] = torch.exp(-0.5 * (rotx ** 2 / sigma_x ** 2 + roty ** 2 / sigma_y ** 2))
+    g /= 2 * np.pi * sigma_x * sigma_y
+    g *= torch.cos(2 * np.pi * frequency * rotx + offset)
+
+    return g
+
+
+def sample_dist(hist, bins, ns, scale='linear'):
+    rand_sample = np.random.rand(ns)
+    if scale == 'linear':
+        rand_sample = np.interp(rand_sample, np.hstack(([0], hist.cumsum())), bins)
+    elif scale == 'log2':
+        rand_sample = np.interp(rand_sample, np.hstack(([0], hist.cumsum())), np.log2(bins))
+        rand_sample = 2**rand_sample
+    elif scale == 'log10':
+        rand_sample = np.interp(rand_sample, np.hstack(([0], hist.cumsum())), np.log10(bins))
+        rand_sample = 10**rand_sample
+    return rand_sample
+