mirror of
https://github.com/wassname/DeepRL.git
synced 2026-08-22 11:40:47 +08:00
67 lines
2.2 KiB
Python
67 lines
2.2 KiB
Python
#######################################################################
|
|
# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com) #
|
|
# Permission given to modify the code as long as you keep this #
|
|
# declaration at the top #
|
|
#######################################################################
|
|
|
|
import torch
|
|
from torch.autograd import Variable
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
|
|
class FullyConnectedNet(nn.Module):
|
|
def __init__(self, dims, optimizer_fn=None, gpu=True):
|
|
super(FullyConnectedNet, self).__init__()
|
|
self.fc1 = nn.Linear(dims[0], dims[1])
|
|
self.fc2 = nn.Linear(dims[1], dims[2])
|
|
self.fc3 = nn.Linear(dims[2], dims[3])
|
|
self.criterion = nn.MSELoss()
|
|
if optimizer_fn is not None:
|
|
self.optimizer = optimizer_fn(self.parameters())
|
|
self.gpu = gpu and torch.cuda.is_available()
|
|
if self.gpu:
|
|
print 'Transferring network to GPU...'
|
|
self.cuda()
|
|
print 'Network transferred.'
|
|
|
|
def forward(self, x):
|
|
x = torch.from_numpy(np.asarray(x, dtype='float32'))
|
|
if self.gpu:
|
|
x = x.cuda()
|
|
x = Variable(x)
|
|
|
|
y = F.relu(self.fc1(x))
|
|
y = F.relu(self.fc2(y))
|
|
y = self.fc3(y)
|
|
return y
|
|
|
|
def sync_with(self, src_net):
|
|
for param_dst, param_src in zip(self.parameters(), src_net.parameters()):
|
|
param_dst.data.copy_(param_src.data)
|
|
|
|
def predict(self, x):
|
|
return self.forward(x).cpu().data.numpy()
|
|
|
|
def learn(self, x, target):
|
|
target = torch.from_numpy(target)
|
|
if self.gpu:
|
|
target = target.cuda()
|
|
target = Variable(target)
|
|
y = self.forward(x)
|
|
loss = self.criterion(y, target)
|
|
self.zero_grad()
|
|
loss.backward()
|
|
self.optimizer.step()
|
|
|
|
def gradient(self, x, actions, rewards):
|
|
y = self.forward(x)
|
|
target = np.copy(y.data.numpy())
|
|
target[np.arange(target.shape[0]), actions] = np.asarray(rewards)
|
|
target = Variable(torch.from_numpy(target))
|
|
loss = self.criterion(y, target)
|
|
loss.backward()
|
|
|
|
def output_transfer(self, y):
|
|
return y
|