brainscore-vision 2.2.4__py3-none-any.whl → 2.2.6__py3-none-any.whl

brainscore_vision/models/evresnet_50_4/evnet/params.py

@@ -0,0 +1,326 @@
+import torch
+import numpy as np
+import scipy.stats as stats
+from .utils import sample_dist
+from typing import Literal
+
+image_size = 224
+visual_degrees = 7
+kernel_size = {'p': 21, 'm': 65}  # 95% of Gaussian coverage
+
+
+# Receptive fields of P and M ganglion cells across the primate retina (Kroner and Kaplan, 1994)
+# https://www.sciencedirect.com/science/article/pii/0042698994E0066T
+
+# P Cells (eccentricity range of 0-5)
+P_cell_params = {
+    'med_rc': 0.03, 'iqr_rc': 0.01,  # Center radius
+    'med_kc': 325.2, 'iqr_kc': 302,  # Center peak sensitivity
+    'med_rs': 0.18, 'iqr_rs': 0.07,  # Surround radius
+    'med_ks': 4.4, 'iqr_ks': 4.6,  # Surround peak sensitivity
+    'c_kc': 0.391, 'm_kc': -1.850,  # Center peak sensitivity vs. radius regression
+    'c_ks': 0.128, 'm_ks': -2.147,  # Surround peak sensitivity vs. radius regression
+}
+
+# M Cells (eccentricity range of 0-10)
+M_cell_params = {
+    'med_rc': 0.10, 'iqr_rc': 0.02,  # Center radius
+    'med_kc': 148.0, 'iqr_kc': 122.4,  # Center peak sensitivity
+    'med_rs': 0.72, 'iqr_rs': 0.23,  # Surround radius
+    'med_ks': 1.1, 'iqr_ks': 0.8,  # Surround peak sensitivity
+}
+
+
+def get_dog_params(
+    features:int, sampling:Literal['median', 'binning', 'uniform', 'lognormal']='median',
+    colors:list[Literal['r/g', 'g/r', 'b/y', 'w/b']]=['r/g', 'g/r', 'b/y'],
+    polarity:list[Literal[0, 1]]=None,
+    cell_type:Literal['p', 'm']='p', image_size:int=image_size, visual_degrees:int=visual_degrees
+    ) -> dict:
+    """Generates DoG parameters for RetinaBlock with more than 3 channels.
+    Number of channels = number of features * 3 color options (R/G, G/R, B/Y).
+    Only generates ON-center cells.
+
+    Args:
+        features (int): _description_
+        binning (bool, optional): whether to use discrete binning while sampling values. Defaults to True.
+        image_size (int, optional): model image size. Defaults to image_size.
+        visual_degrees (int, optional): visual degrees of the model FoV. Defaults to visual_degrees.
+
+    Returns:
+        dict: dictionary with center and surround radii, opponency tensor and DoG kernel size
+    """
+    
+    if not features:
+        return {
+                f'rc_{cell_type}_cell': torch.tensor([]),
+                f'rs_{cell_type}_cell': torch.tensor([]),
+                f'opponency_{cell_type}_cell': torch.tensor([]),
+                f'kernel_{cell_type}_cell': torch.tensor([])
+                }
+
+    assert cell_type in ['p', 'm']
+
+    cell_params = M_cell_params if cell_type=='m' else P_cell_params
+    min_rc = cell_params['med_rc'] - cell_params['iqr_rc']
+    max_rc = cell_params['med_rc'] + cell_params['iqr_rc']
+    min_rs = cell_params['med_rs'] - cell_params['iqr_rs']
+    max_rs = cell_params['med_rs'] + cell_params['iqr_rs']
+
+    color_mapping = {
+            'r/g': np.array([[1,0,0],[0,-1,0]], dtype=np.float16),  # R+/G- (center/surround)
+            'g/r': np.array([[0,1,0],[-1,0,0]], dtype=np.float16),  # G+/R-
+            'b/y': np.array([[0,0,1],[-.5,-.5,0]], dtype=np.float16),  # B+/Y-
+            'w/b': np.array([[1/3]*3,[-1/3]*3], dtype=np.float16)  # ON/OFF
+        }
+    
+    assert features % len(colors) == 0
+    
+    if sampling=='median':
+        # Use median values from distributions (deterministic)
+        assert features == len(colors)
+        rc = np.ones((features,), dtype=np.float16) * cell_params['med_rc']
+        rs = np.ones((features,), dtype=np.float16) * cell_params['med_rs']
+        kc = np.ones((features,), dtype=np.float16) * cell_params['med_kc']
+        ks = np.ones((features,), dtype=np.float16) * cell_params['med_ks']
+    elif sampling=='binning':
+        # Assume uniform joint distribution of rc and rs with discrete binning (deterministic)
+        assert int(np.sqrt(features//len(colors)))==np.sqrt(features//len(colors))
+        edges_rc = np.linspace(min_rc, max_rc, int(np.sqrt(features // len(colors))) + 1)
+        edges_rs = np.linspace(min_rs, max_rs, int(np.sqrt(features // len(colors))) + 1)
+        centers_rc = (edges_rc[:-1] + edges_rc[1:]) / 2
+        centers_rs = (edges_rs[:-1] + edges_rs[1:]) / 2
+        rc = np.repeat(centers_rc, int(np.sqrt(features // len(colors))))
+        rs = np.tile(centers_rs, int(np.sqrt(features // len(colors))))
+    elif sampling=='uniform':
+        # Assume uniform disjoint distribution of rc and rs without binning (stochastic)
+        rc = np.random.uniform(min_rc, max_rc, features // len(colors))
+        rs = np.random.uniform(min_rs, max_rs, features // len(colors))
+    elif sampling=='lognormal':
+        # Assume lognormal disjoint distribution of rc and rs (stochastic)
+        std_rc = (np.log(cell_params['med_rc'] - (cell_params['iqr_rc']/2)) - np.log(cell_params['med_rc'])) / stats.norm.ppf(.25)
+        std_rs = (np.log(cell_params['med_rs'] - (cell_params['iqr_rs']/2)) - np.log(cell_params['med_rs'])) / stats.norm.ppf(.25)
+        rc = np.random.lognormal(np.log(cell_params['med_rc']), std_rc, features // len(colors))
+        rs = np.random.lognormal(np.log(cell_params['med_rs']), std_rs, features // len(colors))
+    
+    if sampling != 'median':
+        assert cell_type == 'p'
+        rc = np.tile(rc, len(colors))
+        rs = np.tile(rs, len(colors))
+        kc = cell_params['c_kc'] * rc ** cell_params['m_kc']
+        ks = cell_params['c_ks'] * rs ** cell_params['m_ks']
+
+    opponency = np.concatenate([
+        np.repeat(color_mapping[c][None, ...], features // len(colors), axis=0)
+        for c in colors
+        ])
+
+    # Conversions
+    ppd = image_size / visual_degrees  # pixels per FOV degree
+    rc = torch.from_numpy(rc * ppd)
+    rs = torch.from_numpy(rs * ppd)
+    kc = torch.from_numpy(kc / ppd ** 2)
+    ks = torch.from_numpy(ks / ppd ** 2)
+    opponency = torch.from_numpy(opponency)
+
+    opponency[:,1] *= torch.unsqueeze(ks[:]/kc[:], 1)
+    
+    if polarity:
+        assert len(polarity) == opponency.size(0)
+        opponency *= torch.tensor(polarity)[..., None, None]
+    
+    params = {
+        f'rc_{cell_type}_cell': rc,
+        f'rs_{cell_type}_cell': rs,
+        f'opponency_{cell_type}_cell': opponency,
+        f'kernel_{cell_type}_cell': kernel_size[cell_type]
+    }
+
+    return params
+
+def get_div_norm_params(
+    relative_size_la, kernel_la=None, 
+    image_size=image_size, visual_degrees=visual_degrees
+    ) -> dict:
+
+    # Conversions
+    ppd = image_size / visual_degrees  # pixels per FOV degree
+    radius_la = P_cell_params['med_rs'] * relative_size_la * ppd
+    radius_cn = 2 * P_cell_params['med_rc'] * ppd
+    c50 = .3
+
+    if not kernel_la and radius_la < np.inf:
+        kernel_la = int(radius_la*2) + int(int(radius_la*2)%2==0)
+    
+    params = {
+        'kernel_la': kernel_la,
+        'radius_la': radius_la,
+        'kernel_cn': kernel_size['p'],
+        'radius_cn': radius_cn,
+        'c50': c50
+    }
+
+    return params
+
+
+def get_grating_params(
+    sf, angle=0, phase=0, contrast=1, radius=.5,
+    image_size=image_size, visual_degrees=visual_degrees
+    ) -> dict:
+    ppd = image_size / visual_degrees  # pixels per FOV degree
+    params = {
+        'size': image_size,
+        'radius': radius * ppd,
+        'sf': sf/ppd,
+        'theta': angle,
+        'phase': phase,
+        'contrast': contrast
+    }
+    return params
+
+
+def generate_gabor_param(
+    n_sc, n_cc, seed=0, rand_flag=False, sf_corr=0.75,
+    sf_max=11.5, sf_min=0, diff_n=False, dnstd=0.22,
+    # Additional parameters
+    in_channels=3, set_orientation=None
+    ):
+    
+    features = n_sc + n_cc
+
+    # 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])
+        nx_dist = np.array([1])
+
+        ny_bins = np.array([0.1, 10**0])
+        ny_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, 11.2])
+        sf_s_dist = np.array([1,  1,  1, 1, 1, 1, 1, 1, 1])
+        sf_c_dist = np.array([1,  1,  1, 1, 1, 1, 1, 1, 1])
+
+    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()
+
+        # Ringach 2002b
+        nx_bins = np.logspace(-1, 0., 5, base=10)
+        ny_bins = np.logspace(-1, 0., 5, base=10)
+        n_joint_dist = np.array([[2.,  0.,  1.,  0.],
+                                 [8.,  9.,  4.,  1.],
+                                 [1.,  2., 19., 17.],
+                                 [0.,  0.,  1.,  7.]])
+        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, 11.2])
+        # foveal only
+        sf_s_dist = np.array([4, 4, 8, 25, 33, 26, 28, 12, 8])
+        sf_c_dist = np.array([0, 0, 9, 9, 7, 10, 23, 12, 14])
+
+    phase = sample_dist(phase_dist, phase_bins, features)
+
+    if set_orientation or set_orientation == 0:
+        ori = np.ones((features,)) * set_orientation
+    else:
+        ori = sample_dist(ori_dist, ori_bins, features)
+    
+    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_s_dist = sf_s_dist[sfmin_ind:sfmax_ind]
+    sf_c_dist = sf_c_dist[sfmin_ind:sfmax_ind]
+
+    sf_s_dist = sf_s_dist / sf_s_dist.sum()
+    sf_c_dist = sf_c_dist / sf_c_dist.sum()
+
+    cov_mat = np.array([[1, sf_corr], [sf_corr, 1]])
+
+    if rand_flag:   # Uniform
+        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
+
+        if diff_n: 
+            ny = sample_dist(ny_dist, ny_bins, features, scale='log10')
+        else:
+            ny = 10**(np.random.normal(np.log10(nx), dnstd))
+            ny[ny<0.1] = 0.1
+            ny[ny>1] = 1
+            # ny = nx
+
+        sf = np.interp(samps_cdf[:,1], np.hstack(([0], sf_s_dist.cumsum())), np.log2(sf_bins))
+        sf = 2**sf
+
+    else:   # Biological
+
+        if n_sc > 0:
+            samps = np.random.multivariate_normal([0, 0], cov_mat, n_sc)
+            samps_cdf = stats.norm.cdf(samps)
+
+            nx_s = np.interp(samps_cdf[:,0], np.hstack(([0], nx_dist.cumsum())), np.log10(nx_bins))
+            nx_s = 10**nx_s
+
+            ny_samp = np.random.rand(n_sc)
+            ny_s = np.zeros(n_sc)
+            for samp_ind, nx_samp in enumerate(nx_s):
+                bin_id = np.argwhere(nx_bins < nx_samp)[-1]
+                ny_s[samp_ind] = np.interp(ny_samp[samp_ind], np.hstack(([0], ny_dist_marg[bin_id, :].cumsum())),
+                                                 np.log10(ny_bins))
+            ny_s = 10**ny_s
+
+            sf_s = np.interp(samps_cdf[:,1], np.hstack(([0], sf_s_dist.cumsum())), np.log2(sf_bins))
+            sf_s = 2**sf_s
+        else:
+            nx_s = np.array([])
+            ny_s = np.array([])
+            sf_s = np.array([])
+
+        if n_cc > 0:
+            samps = np.random.multivariate_normal([0, 0], cov_mat, n_cc)
+            samps_cdf = stats.norm.cdf(samps)
+
+            nx_c = np.interp(samps_cdf[:,0], np.hstack(([0], nx_dist.cumsum())), np.log10(nx_bins))
+            nx_c = 10**nx_c
+
+            ny_samp = np.random.rand(n_cc)
+            ny_c = np.zeros(n_cc)
+            for samp_ind, nx_samp in enumerate(nx_c):
+                bin_id = np.argwhere(nx_bins < nx_samp)[-1]
+                ny_c[samp_ind] = np.interp(ny_samp[samp_ind], np.hstack(([0], ny_dist_marg[bin_id, :].cumsum())),
+                                                 np.log10(ny_bins))
+            ny_c = 10**ny_c
+
+            sf_c = np.interp(samps_cdf[:,1], np.hstack(([0], sf_c_dist.cumsum())), np.log2(sf_bins))
+            sf_c = 2**sf_c
+        else:
+            nx_c = np.array([])
+            ny_c = np.array([])
+            sf_c = np.array([])
+
+        nx = np.concatenate((nx_s, nx_c))
+        ny = np.concatenate((ny_s, ny_c))
+        sf = np.concatenate((sf_s, sf_c))
+
+    color = np.random.randint(low=0, high=in_channels, size=features, dtype=np.int8)
+
+    return sf, ori, phase, nx, ny, color

brainscore_vision/models/evresnet_50_4/evnet/utils.py

@@ -0,0 +1,142 @@
+import math
+import torch
+import numpy as np
+import random
+
+def gaussian_kernel(
+    sigma: float, k: float=1, size:float=15, norm:bool=False
+    ) -> torch.tensor:
+    """Returns a 2D Gaussian kernel.
+
+    :param sigma (float): standard deviation of the Gaussian
+    :param k (float, optional): height of the Gaussian
+    :param size (float, optional): kernel size
+    :param norm (bool, optional): whether no normalize the kernel
+    :return: gaussian kernel
+    """
+    assert size % 2 == 1
+    w = size // 2
+    grid_val = torch.arange(-w, w+1, dtype=torch.float)
+    x, y = torch.meshgrid(grid_val, grid_val, indexing='ij')
+    gaussian = k * torch.exp(-(x**2 + y**2) / (2*(sigma)**2))
+    if norm: gaussian /= torch.abs(gaussian.sum())
+    return gaussian
+
+
+def dog_kernel(
+    sigma_c: float, sigma_s: float, k_c: float, k_s: float,
+    polarity:int, size:int=21
+    ) -> torch.tensor:
+    """Returns a 2D Difference-of-Gaussians kernel.
+
+    :param sigma_c: standard deviation of the center Gaussian
+    :param sigma_s: standard deviation of the surround Gaussian
+    :param k_c: peak sensitivity of the center
+    :param k_s: peak sensitivity of the surround
+    :param polarity: polarity of the center Gaussian (+1 or -1)
+    :param size: kernel size
+    :return: difference-of-gaussians kernel
+    """
+    assert size % 2 == 1
+    assert polarity in [-1 , 1]
+    center_gaussian = gaussian_kernel(sigma=sigma_c, k=k_c, size=size)
+    surround_gaussian = gaussian_kernel(sigma=sigma_s, k=k_s, size=size)
+    dog = polarity * (center_gaussian - surround_gaussian)
+    dog /= torch.sum(dog)
+    return dog
+
+def circular_kernel(size:int, radius:float) -> torch.tensor:
+    """Returns circular kernel.
+
+    :param size (int): kernel size
+    :param radius (float): radius of the circle
+    :return: circular kernel
+    """
+
+    w = size // 2
+    grid_val = torch.arange(-w, w+1, dtype=torch.float)
+    x, y = torch.meshgrid(grid_val, grid_val, indexing='ij')
+    kernel = torch.zeros(y.shape)
+    kernel[torch.sqrt(x**2 + y**2) <= radius] = 1
+    kernel /= torch.sum(kernel)
+    return kernel
+
+def gabor_kernel(
+    frequency:float,  sigma_x:float, sigma_y:float,
+    theta:float=0, offset:float=0, ks:int=61
+    ):
+    """Returns gabor kernel.
+
+    :param frequency (float): spatial frequency of gabor
+    :param sigma_x (float): standard deviation in x direction
+    :param sigma_y (float): standard deviation in y direction
+    :param theta (int, optional): Angle theta. Defaults to 0.
+    :param offset (int, optional): Offset. Defaults to 0.
+    :param ks (int, optional): Kernel size. Defaults to 61.
+    :return: np.ndarray: 2-dimensional Gabor kernel
+    """
+    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 generate_grating(
+    size:int, radius:float, sf:float, theta:float=0, phase:float=0,
+    contrast:float=1, gaussian_mask:bool=False
+    ) -> torch.tensor:
+    """Returns masked grating array.
+
+    :param size (int): kernel size
+    :param radius (float): standard deviation times sqrt(2) of the mask if gaussian_mask is True, and the radius if is false
+    :param sf (float): spatial frequency of the grating
+    :param theta (float, optional): angle of the grating 
+    :param phase (float, optional): phase of the grating
+    :param gaussian_mask (bool, optional): mask is a Gaussian if true and a circle if false 
+    :param contrast (float, optional): maximum contrast of the grating
+    :return: 2d masked grating array
+    """
+    grid_val = torch.linspace(-size//2, size//2+1, size, dtype=torch.float)
+    X, Y = torch.meshgrid(grid_val, grid_val, indexing='ij')
+    grating = torch.sin(2*math.pi*sf*(X*math.cos(theta) + Y*math.sin(theta)) + phase) * contrast
+    mask = torch.exp(-((X**2 + Y**2)/(2*(radius/np.sqrt(2))**2))) if gaussian_mask else torch.sqrt(X**2 + Y**2) <= radius
+    return grating * mask * .5 + .5
+
+
+def sample_dist(hist:np.array, bins:int, ns:float, scale:str='linear'):
+    """Samples from distributions with different scales.
+
+    Args:
+        hist (np.array): histogram
+        bins (int): number of bins
+        ns (float): sample size
+        scale (str, optional): distribution scale. Defaults to 'linear'.
+
+    :returns rand_sample (np.array): 
+    """    
+    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
+
+def set_seed(seed):
+    """Enforces deterministic behaviour and sets RNG seed for numpy and pytorch.
+
+    :param seed (int): seed
+    """
+    random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+

brainscore_vision/models/evresnet_50_4/model.py

@@ -0,0 +1,67 @@
+import functools
+from brainscore_vision.model_helpers.activations.pytorch import PytorchWrapper
+from brainscore_vision.model_helpers.check_submission import check_models
+from brainscore_vision.model_helpers.activations.pytorch import load_preprocess_images
+from brainscore_vision.model_helpers.s3 import load_weight_file
+import torch
+from .evnet.evnet import EVNet, evnet_params
+
+def get_model(name):
+    assert name == 'evresnet_50_4'
+    model = EVNet(
+        **evnet_params['evnet'], model_arch='resnet50', image_size=224,
+        gabor_seed=1, visual_degrees=7, num_classes=1000,
+        sf_max=11.5, k_exc=25, stride=4
+        )
+    weight_file = load_weight_file(
+        bucket="evnets-model-weights",
+        relative_path="evresnet_50_1.pth",
+        sha1="ca92ffc54171dd0eee1576b6ff375d6b02623aa7",
+        version_id="LuiUBSKSo8r0qVAi8UuaFCBVoQxszmWw"
+        )
+    model.to(torch.device('cpu'))
+    checkpoint = torch.load(weight_file, map_location=torch.device('cpu'))
+    model.load_state_dict(checkpoint['model'])
+    preprocessing = functools.partial(
+        load_preprocess_images,
+        image_size=224,
+        normalize_mean=(0,0,0),
+        normalize_std=(1,1,1)
+        )
+    wrapper = PytorchWrapper(
+        identifier='evresnet_50_4',
+        model=model, preprocessing=preprocessing
+        )
+    wrapper.image_size = 224
+    return wrapper
+
+
+def get_layers(name):
+    assert name == 'evresnet_50_4'
+    layers = (
+            ['voneblock'] +
+            ['model.0.0', 'model.0.1', 'model.0.2'] +
+            ['model.1.0', 'model.1.1', 'model.1.2', 'model.1.3'] +
+            ['model.2.0', 'model.2.1', 'model.2.2', 'model.2.3',
+             'model.2.4', 'model.2.5'] +
+            ['model.3.0', 'model.3.1', 'model.3.2'] +
+            ['model.5']
+    )
+    return layers
+
+
+def get_bibtex(model_identifier):
+    return """@misc{piper2024explicitlymodelingprecorticalvision,
+      title={Explicitly Modeling Pre-Cortical Vision with a Neuro-Inspired Front-End Improves CNN Robustness}, 
+      author={Lucas Piper and Arlindo L. Oliveira and Tiago Marques},
+      year={2024},
+      eprint={2409.16838},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2409.16838}
+      }"""
+
+
+if __name__ == '__main__':
+    get_model('evresnet_50_4')
+    check_models.check_base_models(__name__)

brainscore_vision/models/evresnet_50_4/requirements.txt

@@ -0,0 +1,4 @@
+torch
+torchvision
+numpy
+scipy

brainscore_vision/models/evresnet_50_4/test.py

@@ -0,0 +1,8 @@
+import pytest
+import brainscore_vision
+
+
+@pytest.mark.travis_slow
+def test_has_identifier():
+    model = brainscore_vision.load_model('evresnet_50_4')
+    assert model.identifier == 'evresnet_50_4'

brainscore_vision/models/evresnet_50_4_no_mapping/__init__.py

@@ -0,0 +1,10 @@
+from brainscore_vision import model_registry
+from brainscore_vision.model_helpers.brain_transformation import ModelCommitment
+from .model import get_model, get_layers
+
+model_registry['evresnet_50_4_no_mapping'] = lambda: ModelCommitment(
+    identifier='evresnet_50_4_no_mapping',
+    activations_model=get_model('evresnet_50_4_no_mapping'),
+    layers=get_layers('evresnet_50_4_no_mapping'),
+    visual_degrees=7
+    )

brainscore_vision/models/evresnet_50_4_no_mapping/evnet/backends.py

@@ -0,0 +1,109 @@
+import torch
+from torch import nn
+import torchvision
+
+def get_resnet_backend(
+    num_classes:int,
+    p_channels:int,
+    m_channels:int,
+    layers:int=18,
+    with_voneblock:bool=False,
+    tiny:bool=False
+    ):
+    """Returns ResNet18 or ResNet50 backend.
+
+    :param num_classes (int): number of classes in the classifier
+    :param p_channels (int, optional): number of Midget/P cell channels
+    :param m_channels (int, optional): number of Parasol/M cell channels
+    :param layers (int, optional): number of architecture layers
+    :param with_voneblock (bool, optional): whether to remove the first block of the backend
+    :param tiny (bool, optional): whether to employ Tiny ImageNet adaptation (64px/2deg input)
+    :return: backend model, number of backend in channels
+    """
+    assert layers in [18, 50]
+    backend = torchvision.models.resnet18() if layers == 18 else torchvision.models.resnet50()
+    backend.fc = nn.Linear(backend.fc.in_features, num_classes)
+    if with_voneblock:
+        # When using the VOneBlock, in channels are defined in the bottleneck
+        backend_in_channels = backend.layer1[0].conv1.in_channels
+        backend = nn.Sequential(
+                *list(backend.children())[4:-1],  # Remove first block from ResNet-18
+                nn.Flatten(),
+                backend.fc
+                )
+    else:
+        backend_in_channels = 3
+        in_channels = p_channels + m_channels
+        if tiny: backend.conv1.stride = (1, 1)
+        conv1 = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=backend.conv1.out_channels,
+            kernel_size=backend.conv1.kernel_size,
+            stride=backend.conv1.stride,
+            padding=backend.conv1.padding,
+            bias=backend.conv1.bias
+            )
+        weight = torch.zeros_like(conv1.weight.data)
+        weight[:, :min(in_channels, backend_in_channels), :, :] =\
+            backend.conv1.weight.data[:, :min(in_channels, backend_in_channels), :, :]
+        if weight.size(1) > backend.conv1.weight.data.size(1):
+            nn.init.kaiming_normal_(weight[:, in_channels:, :, :], mode="fan_out", nonlinearity="relu")
+        conv1.weight.data = weight
+        backend.conv1 = conv1
+    
+    return backend, backend_in_channels
+
+
+def get_vgg_backend(
+    num_classes:int,
+    p_channels:int,
+    m_channels:int,
+    layers:int=16,
+    with_voneblock:bool=False,
+    tiny:bool=False
+    ):
+    """Returns VGG16 or VGG19 backend.
+
+    :param num_classes (int): number of classes in the classifier
+    :param p_channels (float, optional): number of Midget/P cell channels
+    :param m_channels (float, optional): number of Parasol/M cell channels
+    :param layers (int, optional): number of architecture layers
+    :param with_voneblock (bool, optional): whether to remove the first block of the backend
+    :param tiny (bool, optional): whether to employ Tiny ImageNet adaptation (64px/2deg input)
+    :return: backend model, number of backend in channels
+    """
+    assert layers in [16, 19]
+    backend = torchvision.models.vgg16() if layers == 16 else torchvision.models.vgg19()
+    backend.classifier[-1] = nn.Linear(backend.classifier[-1].in_features, num_classes)
+    if tiny:
+        backend.features = nn.Sequential(*list(backend.features[:-1]))
+        backend.classifier[0] = nn.Linear(in_features=25088, out_features=2048, bias=True)
+        backend.classifier[3] = nn.Linear(in_features=2048, out_features=2048, bias=True)
+        backend.classifier[6] = nn.Linear(in_features=2048, out_features=200, bias=True)
+    if with_voneblock:
+        backend_in_channels = backend.features[2].in_channels
+        backend.features = nn.Sequential(
+                *list(backend.features[2:])
+                )
+    else:
+        backend_in_channels = 3
+        in_channels = p_channels + m_channels
+        backend.features[0] = nn.Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=True)
+        conv1 = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=backend.features[0].out_channels,
+            kernel_size=backend.features[0].kernel_size,
+            stride=backend.features[0].stride,
+            padding=backend.features[0].padding,
+            bias=True
+            )
+        weight = torch.zeros_like(conv1.weight.data)
+        weight[:, :min(in_channels, backend_in_channels), :, :] =\
+            backend.conv1.weight.data[:, :min(in_channels, backend_in_channels), :, :]
+        if weight.size(1) > backend.conv1.weight.data.size(1):
+            nn.init.kaiming_normal_(weight[:, in_channels:, :, :], mode="fan_out", nonlinearity="relu")
+        conv1.bias.data = backend.features[0].bias.data
+        conv1.weight.data = weight
+        backend.features[0] = conv1
+    
+    return backend, backend_in_channels