def __init__(self,
in_dim,
hid_dim,
out_dim,
hid_layers_activation=None,
optim_param=None,
loss_param=None,
clamp_grad=False,
clamp_grad_val=1.0):
'''
in_dim: dimension of the inputs
hid_dim: list containing dimensions of the hidden layers
out_dim: list containing the dimensions of the ouputs
hid_layers_activation: activation function for the hidden layers
optim_param: parameters for initializing the optimizer
loss_param: measure of error between model predictions and correct outputs
clamp_grad: whether to clamp the gradient
@example:
net = MLPHeterogenousHeads(
1000,
[512, 256, 128],
[1, 1],
hid_layers_activation='relu',
optim_param={'name': 'Adam'},
loss_param={'name': 'mse_loss'},
clamp_grad=True,
clamp_grad_val=2.0)
'''
nn.Module.__init__(self)
# Create net and initialize params
self.in_dim = in_dim
self.out_dim = out_dim
self.layers = []
# Init network body
for i, layer in enumerate(hid_dim):
in_D = in_dim if i == 0 else hid_dim[i - 1]
out_D = hid_dim[i]
self.layers += [nn.Linear(in_D, out_D)]
self.layers += [net_util.get_activation_fn(hid_layers_activation)]
in_D = hid_dim[-1] if len(hid_dim) > 0 else in_dim
self.body = nn.Sequential(*self.layers)
# Init network output heads
self.out_layers = []
for i, dim in enumerate(out_dim):
self.out_layers += [nn.Linear(in_D, dim)]
self.layers += [self.out_layers]
self.init_params()
# Init other net variables
self.params = list(self.body.parameters())
for layer in self.out_layers:
self.params.extend(list(layer.parameters()))
self.optim_param = optim_param
self.optim = net_util.get_optim_multinet(self.params, self.optim_param)
self.loss_fn = net_util.get_loss_fn(self, loss_param)
self.clamp_grad = clamp_grad
self.clamp_grad_val = clamp_grad_val
def make_shared_body(self, in_dim, dims, hid_layers_activation):
'''Creates the shared body of the network. Stored as a Sequential
model in self.body'''
for i, layer in enumerate(dims):
in_D = in_dim if i == 0 else dims[i - 1]
out_D = dims[i]
self.shared_layers += [nn.Linear(in_D, out_D)]
self.shared_layers += [
net_util.get_activation_fn(hid_layers_activation)
]
return nn.Sequential(*self.shared_layers)
def __init__(self,
in_dim,
hid_dim,
out_dim,
hid_layers_activation=None,
optim_param=None,
loss_param=None,
clamp_grad=False,
clamp_grad_val=1.0,
gpu=False):
'''
in_dim: dimension of the inputs
hid_dim: list containing dimensions of the hidden layers
out_dim: dimension of the ouputs
hid_layers_activation: activation function for the hidden layers
optim_param: parameters for initializing the optimizer
loss_param: measure of error between model predictions and correct outputs
clamp_grad: whether to clamp the gradient
gpu: whether to train using a GPU. Note this will only work if a GPU is available, othewise setting gpu=True does nothing
@example:
net = MLPNet(
1000,
[512, 256, 128],
10,
hid_layers_activation='relu',
optim_param={'name': 'Adam'},
loss_param={'name': 'mse_loss'},
clamp_grad=True,
clamp_grad_val=2.0,
gpu=True)
'''
super(MLPNet, self).__init__()
# Create net and initialize params
self.in_dim = in_dim
self.out_dim = out_dim
self.layers = []
for i, layer in enumerate(hid_dim):
in_D = in_dim if i == 0 else hid_dim[i - 1]
out_D = hid_dim[i]
self.layers += [nn.Linear(in_D, out_D)]
self.layers += [net_util.get_activation_fn(hid_layers_activation)]
in_D = hid_dim[-1] if len(hid_dim) > 0 else in_dim
self.layers += [nn.Linear(in_D, out_dim)]
self.model = nn.Sequential(*self.layers)
self.init_params()
if torch.cuda.is_available() and gpu:
self.model.cuda()
# Init other net variables
self.params = list(self.model.parameters())
self.optim_param = optim_param
self.optim = net_util.get_optim(self, self.optim_param)
self.loss_fn = net_util.get_loss_fn(self, loss_param)
self.clamp_grad = clamp_grad
self.clamp_grad_val = clamp_grad_val
def build_state_proc_layers(self, state_processing_layers, hid_layers_activation):
'''Builds all of the state processing layers in the network.
These layers are turned into a Sequential model and stored
in self.state_proc_model.
The entire model consists of three parts:
1. self.state_proc_model
2. self.rnn
3. self.fc_out'''
for i, layer in enumerate(state_processing_layers):
in_D = self.in_dim if i == 0 else state_processing_layers[i - 1]
out_D = state_processing_layers[i]
self.state_processing_layers += [nn.Linear(in_D, out_D)]
self.state_processing_layers += [net_util.get_activation_fn(hid_layers_activation)]
return nn.Sequential(*self.state_processing_layers)
def build_flat_layers(self, flat_hid, out_dim, hid_layers_activation):
'''Builds all of the dense layers in the network.
These layers are turned into a Sequential model and stored
in self.dense_model.
The entire model consists of two parts:
1. self.conv_model
2. self.dense_model'''
self.flat_dim = self.get_conv_output_size()
for i, layer in enumerate(flat_hid):
in_D = self.flat_dim if i == 0 else flat_hid[i - 1]
out_D = flat_hid[i]
self.flat_layers += [nn.Linear(in_D, out_D)]
self.flat_layers += [
net_util.get_activation_fn(hid_layers_activation)]
in_D = flat_hid[-1] if len(flat_hid) > 0 else self.flat_dim
self.flat_layers += [nn.Linear(in_D, out_dim)]
return nn.Sequential(*self.flat_layers)
def build_conv_layers(self, conv_hid_layers):
'''
Builds all of the convolutional layers in the network and store in a Sequential model
'''
conv_layers = []
in_d = 1 # input channel
for i, hid_layer in enumerate(conv_hid_layers):
hid_layer = [tuple(e) if ps.is_list(e) else e for e in hid_layer] # guard list-to-tuple
# hid_layer = out_d, kernel, stride, padding, dilation
conv_layers.append(nn.Conv2d(in_d, *hid_layer))
conv_layers.append(net_util.get_activation_fn(self.hid_layers_activation))
# Don't include batch norm in the first layer
if self.batch_norm and i != 0:
conv_layers.append(nn.BatchNorm2d(in_d))
in_d = hid_layer[0] # update to out_d
conv_model = nn.Sequential(*conv_layers)
return conv_model
def build_conv_layers(self, conv_hid_layers):
'''
Builds all of the convolutional layers in the network and store in a Sequential model
'''
conv_layers = []
for i, hid_layer in enumerate(conv_hid_layers):
conv_layers.append(nn.Conv2d(
hid_layer[0], # in chnl
hid_layer[1], # out chnl
tuple(hid_layer[2]), # kernel
stride=hid_layer[3],
padding=hid_layer[4],
dilation=tuple(hid_layer[5])))
conv_layers.append(net_util.get_activation_fn(self.hid_layers_activation))
# Don't include batch norm in the first layer
if self.batch_norm and i != 0:
conv_layers.append(nn.BatchNorm2d(hid_layer[1]))
conv_model = nn.Sequential(*conv_layers)
return conv_model
def make_action_heads(self, in_dim, act_heads, hid_layers_activation):
'''Creates each action head. These are stored as Sequential
models in self.action_heads_models'''
act_heads_models = []
for head in act_heads:
layers = []
assert len(head) > 0
for i, layer in enumerate(head):
in_D = head[i - 1] if i > 0 else in_dim
out_D = head[i]
layers += [nn.Linear(in_D, out_D)]
# No activation function in the last layer
if i < len(head) - 1:
layers += [
net_util.get_activation_fn(hid_layers_activation)
]
self.action_heads_layers.append(layers)
act_heads_models.append(nn.Sequential(*layers))
return act_heads_models
def build_conv_layers(self, conv_hid_layers):
'''
Builds all of the convolutional layers in the network and store in a Sequential model
'''
conv_layers = []
for i, hid_layer in enumerate(conv_hid_layers):
conv_layers.append(nn.Conv2d(
hid_layer[0], # in chnl
hid_layer[1], # out chnl
tuple(hid_layer[2]), # kernel
stride=hid_layer[3],
padding=hid_layer[4],
dilation=tuple(hid_layer[5])))
conv_layers.append(net_util.get_activation_fn(self.hid_layers_activation))
# Don't include batch norm in the first layer
if self.batch_norm and i != 0:
conv_layers.append(nn.BatchNorm2d(hid_layer[1]))
conv_model = nn.Sequential(*conv_layers)
return conv_model
def make_state_heads(self, state_heads, hid_layers_activation):
'''Creates each state head. These are stored as Sequential
models in self.state_heads_models'''
self.state_out_concat = 0
state_heads_models = []
for head in state_heads:
layers = []
assert len(head) > 1
for i, layer in enumerate(head):
if i != 0:
in_D = head[i - 1]
out_D = head[i]
layers += [nn.Linear(in_D, out_D)]
layers += [
net_util.get_activation_fn(hid_layers_activation)
]
self.state_out_concat += head[-1]
self.state_heads_layers.append(layers)
state_heads_models.append(nn.Sequential(*layers))
return state_heads_models
def build_conv_layers(self, conv_hid, hid_layers_activation):
'''Builds all of the convolutional layers in the network.
These layers are turned into a Sequential model and stored
in self.conv_model.
The entire model consists of two parts:
1. self.conv_model
2. self.dense_model'''
for i, layer in enumerate(conv_hid):
self.conv_layers += [nn.Conv2d(
conv_hid[i][0],
conv_hid[i][1],
conv_hid[i][2],
stride=conv_hid[i][3],
padding=conv_hid[i][4],
dilation=conv_hid[i][5])]
self.conv_layers += [
net_util.get_activation_fn(hid_layers_activation)]
# Don't include batch norm in the first layer
if self.batch_norm and i != 0:
self.conv_layers += [nn.BatchNorm2d(conv_hid[i][1])]
return nn.Sequential(*self.conv_layers)
