EmergentPredictiveCoding/functions.py at master · KietzmannLab/EmergentPredictiveCoding · GitHub

History

93 lines (73 loc) · 2.12 KB

Raw

import numpy as np

import torch

import torch.nn.functional as F

import time

from typing import Dict

def get_device():

if torch.cuda.is_available():

DEVICE = 'cuda'

torch.set_default_tensor_type(torch.cuda.FloatTensor)

elif torch.backends.mps.is_available():

DEVICE = torch.device("mps")

else:

DEVICE = 'cpu'

print('Using {}'.format(DEVICE))

return DEVICE

def get_time() -> int:

'''Returns current time in ms'''

return int(round(time.time() * 1000))

class Timer:

def __init__(self):

self.reset()

def lap(self):

self.t_lap = get_time()

def get(self):

return get_time() - self.t_lap

def reset(self):

self.t_total = get_time()

self.t_lap = get_time()

def __str__(self):

t = self.get()

ms = t % 1000

t = int(t / 1000)

s = t % 60

t = int(t / 60)

m = t % 60

if t == 0:

return "{}.{:03}".format(s,ms)

else:

t = int(t / 60)

h = t

if t == 0:

return "{}:{:02}.{:03}".format(m,s,ms)

else:

return "{}:{:02}:{:02}.{:03}".format(h,m,s,ms)

def append_dict(dict_a:Dict[str,np.ndarray], dict_b:Dict[str,np.ndarray]):

for k, v in dict_b.items():

dict_a[k] = np.concatenate((dict_a[k], v))

def L1Loss(x:torch.FloatTensor):

return torch.mean(torch.abs(x))

def L2Loss(x:torch.FloatTensor):

return torch.mean(torch.pow(x, 2))

def Linear(x:torch.FloatTensor):

return x

def parse_loss(args, terms):

if args == None:

return L1Loss, torch.tensor(0.0)

pre, post, weights = terms

arg1, arg2 = args.split('_')

if arg1 == 'l1':

loss_fn = L1Loss

else:

loss_fn = L2Loss

if arg2 == 'pre':

loss_arg = pre

elif arg2 == 'post':

loss_arg = post

else:

loss_arg = weights

return loss_fn, loss_arg

def init_params(size_x, size_y):

return ((torch.rand(size_x, size_y) * 2.) - 1. ) * np.sqrt(1. / size_x)

def normalize(x, p=2.0, dim=1):

return F.normalize(x, p, dim)

EmergentPredictiveCoding/functions.py at master · KietzmannLab/EmergentPredictiveCoding · GitHub

Expand file tree

Breadcrumbs

functions.py

Latest commit

History

Breadcrumbs

functions.py

File metadata and controls