|
- import shutil
- import torch
- import os
- import numpy as np
-
- def get_n_params(model):
- pp=0
- for p in list(model.parameters()):
- nn=1
- for s in list(p.size()):
- nn= nn*s
- pp += nn
- return pp
-
- def save_checkpoint(state,is_best,save_path,filename = 'checkpoint.pth.tar'):
- torch.save(state,os.path.join(save_path,filename))
- if is_best:
- shutil.copyfile(os.path.join(save_path,filename), os.path.join(save_path, 'model_best.pth.tar'))
-
-
-
- class AverageMeter(object):
- """Computes and stores the average and current value"""
-
- def __init__(self):
- self.reset()
-
- def reset(self):
- self.val = 0
- self.avg = 0
- self.sum = 0
- self.count = 0
-
- def update(self, val, n=1):
- self.val = val
- self.sum += val * n
- self.count += n
- self.avg = self.sum / self.count
-
- def __repr__(self):
- return '{:.3f} ({:.3f})'.format(self.val, self.avg)
-
-
- # Simple replay buffer
- class ReplayBuffer(object):
- def __init__(self):
- self.storage = []
-
- # Expects tuples of (state, next_state, action, reward, done)
- def add(self, data):
- self.storage.append(data)
-
- def sample(self, batch_size=100):
- ind = np.random.randint(0, len(self.storage), size=batch_size)
- x, y, u, r, d = [], [], [], [], []
-
- for i in ind:
- X, Y, U, R, D = self.storage[i]
- x.append(np.array(X, copy=False))
- y.append(np.array(Y, copy=False))
- u.append(np.array(U, copy=False))
- r.append(np.array(R, copy=False))
- d.append(np.array(D, copy=False))
-
- return np.array(x), np.array(y), np.array(u), np.array(r).reshape(-1, 1), np.array(d).reshape(-1, 1)
|