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)