mirror of
https://github.com/wassname/DeepRL.git
synced 2026-10-08 11:39:05 +08:00
17 lines
536 B
Python
17 lines
536 B
Python
# adapted from https://github.com/alexis-jacq/Pytorch-DPPO/blob/master/model.py
|
|
|
|
import torch
|
|
|
|
class SharedGrad():
|
|
def __init__(self, model):
|
|
self.grads = {}
|
|
for name, p in model.named_parameters():
|
|
self.grads[name+'_grad'] = torch.ones(p.size()).share_memory_()
|
|
|
|
def add_gradient(self, model):
|
|
for name, p in model.named_parameters():
|
|
self.grads[name+'_grad'] += p.grad.data
|
|
|
|
def reset(self):
|
|
for name,grad in self.grads.items():
|
|
self.grads[name].fill_(0) |