def __init__(self, statistics_group_size=1, name=None, data_layout='data_parallel'): super().__init__() BatchNormModule.global_count += 1 self.instance = 0 self.statistics_group_size = statistics_group_size self.name = (name if name else 'bnmodule{0}'.format( BatchNormModule.global_count)) self.data_layout = data_layout # Initialize weights self.scale = lbann.Weights( initializer=lbann.ConstantInitializer(value=1.0), name=self.name + '_scale') self.bias = lbann.Weights( initializer=lbann.ConstantInitializer(value=0.0), name=self.name + '_bias') self.running_mean = lbann.Weights( initializer=lbann.ConstantInitializer(value=0.0), name=self.name + '_running_mean') self.running_variance = lbann.Weights( initializer=lbann.ConstantInitializer(value=1.0), name=self.name + '_running_variance')
def __init__(self, size, statistics_group_size=1, name=None, data_layout='data_parallel'): super().__init__() FcBnRelu.global_count += 1 self.instance = 0 self.name = (name if name else 'fcbnrelu{0}'.format(FcBnRelu.global_count)) self.data_layout = data_layout self.fc = lbann.modules.FullyConnectedModule( size, bias=False, name=self.name + '_fc', data_layout=self.data_layout) # Weights for batchnorm scalebias_vals = [1.0] * size + [0.0] * size self.bn_weights = [ lbann.Weights(name='{0}_bn_running_mean'.format(self.name), initializer=lbann.ConstantInitializer(value=0.0)), lbann.Weights(name='{0}_bn_running_var'.format(self.name), initializer=lbann.ConstantInitializer(value=1.0)), lbann.Weights(name='{0}_bn_scalebias'.format(self.name), initializer=lbann.ValueInitializer( values=' '.join([str(x) for x in scalebias_vals]))) ]
def __init__(self, embed_dim, num_heads, branches, d_kv=None, name=None): super().__init__() MultiheadAttention.global_count += 1 self.instance = 0 assert embed_dim % num_heads == 0, 'embed_dim must be divisible by num_heads' self.embed_dim = embed_dim self.num_heads = num_heads self.head_dim = embed_dim // num_heads if (d_kv == None): self.inner_dim = embed_dim self.head_dim = embed_dim // num_heads else: self.inner_dim = d_kv * num_heads self.head_dim = d_kv if (branches == 0): self.ENABLE_SUBGRAPH = False self.BRANCHES = 0 else: self.ENABLE_SUBGRAPH = True self.BRANCHES = branches # Module name self.name = name if not self.name: self.name = f'multiheadattention{MultiheadAttention.global_count}' # Weights for fully-connected layers self.query_weights = [ lbann.Weights(initializer=lbann.GlorotNormalInitializer(), name=f'{self.name}_query_matrix'), lbann.Weights(initializer=lbann.ConstantInitializer(value=0), name=f'{self.name}_query_bias'), ] self.key_weights = [ lbann.Weights(initializer=lbann.GlorotNormalInitializer(), name=f'{self.name}_key_matrix'), lbann.Weights(initializer=lbann.ConstantInitializer(value=0), name=f'{self.name}_key_bias'), ] self.value_weights = [ lbann.Weights(initializer=lbann.GlorotNormalInitializer(), name=f'{self.name}_value_matrix'), lbann.Weights(initializer=lbann.ConstantInitializer(value=0), name=f'{self.name}_value_bias'), ] #Channelwise FC in SubGraph self.output_weights = [] for head in range(branches): self.output_weights.append([ lbann.Weights(initializer=lbann.GlorotNormalInitializer(), name=f'{self.name}_head{head}_output_matrix'), lbann.Weights(initializer=lbann.ConstantInitializer(value=0), name=f'{self.name}_head{head}_output_bias'), ])
def __init__(self, out_channels, kernel_size, stride, padding, use_bn, bn_zero_init, bn_statistics_group_size, activation, parallel_strategy, name, conv_weights): """Initialize ConvBNRelu module. Args: out_channels (int): Number of output channels, i.e. number of convolution filters. kernel_size (int): Size of convolution kernel. stride (int): Convolution stride. padding (int): Convolution padding. use_bn (bool): Whether or not batch normalization layers are used. bn_zero_init (bool): Zero-initialize batch normalization scale. bn_statistics_group_size (int): Aggregation size for batch normalization statistics. activation (lbann.Layer): The activation function. name (str): Module name. conv_weights (lbann.Weights): Pre-defined weights. """ super().__init__() self.name = name self.instance = 0 self.stride = stride self.bn_statistics_group_size = bn_statistics_group_size self.activation = activation self.use_bn = use_bn self.conv_weights = conv_weights self.ps = parallel_strategy # Initialize convolution self.conv = lbann.modules.Convolution3dModule( out_channels, kernel_size, stride=1, padding=padding, bias=False, parallel_strategy=self.ps, weights=self.conv_weights, name=self.name + '_conv') # Initialize batch normalization if self.use_bn: bn_scale_init = 0.0 if bn_zero_init else 1.0 bn_scale = lbann.Weights( initializer=lbann.ConstantInitializer(value=bn_scale_init), name=self.name + '_bn_scale') bn_bias = lbann.Weights( initializer=lbann.ConstantInitializer(value=0.0), name=self.name + '_bn_bias') self.bn_weights = [bn_scale, bn_bias]
def __init__(self, *args, **kwargs): super().__init__() self.name = kwargs["name"] self.activation = None if "activation" not in kwargs.keys() \ else kwargs["activation"] kwargs["activation"] = None self.conv = lm.Convolution3dModule(*args, **kwargs) bn_scale = lbann.Weights( initializer=lbann.ConstantInitializer(value=1.0), name="{}_bn_scale".format(self.name)) bn_bias = lbann.Weights( initializer=lbann.ConstantInitializer(value=0.0), name="{}_bn_bias".format(self.name)) self.bn_weights = [bn_scale, bn_bias] self.instance = 0
def forward_discriminator2(self,img): ''' Discriminator 2. Weights are frozen as part of Adversarial network = Stacked G + D ''' bn_wts=[lbann.Weights(initializer=lbann.ConstantInitializer(value=1.0)), lbann.Weights(initializer=lbann.ConstantInitializer(value=0.0))] for count,lyr in enumerate(self.d2_conv): if count==0: x=lbann.LeakyRelu(lyr(img), negative_slope=0.2) else : x = lbann.LeakyRelu(lyr(x), negative_slope=0.2) #### without convbrlu # if count==0: x = lbann.LeakyRelu(lbann.BatchNormalization(lyr(img),weights=bn_wts,statistics_group_size=-1),negative_slope=0.2) # else: x = lbann.LeakyRelu(lbann.BatchNormalization(lyr(x),weights=bn_wts,statistics_group_size=-1),negative_slope=0.2) dims=524288 y= self.d2_fc(lbann.Reshape(x,dims=str(dims))) return y
def __init__(self, size, bias=False, weights=[], activation=None, transpose=False, name=None, parallel_strategy={}): """Initalize channelwise fully connected module Args: size (int or list): Dimension of the output tensor bias (bool): Whether to apply bias after linearity. transpose (bool): Whether to apply transpose of weights matrix. weights (`Weights` or iterator of `Weights`): Weights in fully-connected layer. There are at most two: the matrix and the bias. If weights are not provided, the matrix will be initialized with He normal initialization and the bias with zeros. activation (type): Layer class for activation function. name (str): Default name is in the form 'channelwisefc<index>'. parallel_strategy (dict): Data partitioning scheme. """ super().__init__() ChannelwiseFullyConnectedModule.global_count += 1 self.instance = 0 self.size = size self.bias = bias self.transpose = transpose self.parallel_strategy = parallel_strategy self.name = (name if name else 'channelwisefc{0}'.format( ChannelwiseFullyConnectedModule.global_count)) self.data_layout = 'data_parallel' self.weights = list(make_iterable(weights)) if len(self.weights) > 2: raise ValueError('`FullyConnectedModule` has ' 'at most two weights, ' 'but got {0}'.format(len(self.weights))) if len(self.weights) == 0: self.weights.append( lbann.Weights(initializer=lbann.HeNormalInitializer(), name=self.name + '_matrix')) if self.bias and len(self.weights) == 1: self.weights.append( lbann.Weights(initializer=lbann.ConstantInitializer(value=0.0), name=self.name + '_bias')) self.activation = None if activation: if isinstance(activation, type): self.activation = activation else: self.activation = type(activation) if not issubclass(self.activation, lbann.Layer): raise ValueError('activation must be a layer')
def __init__(self, out_channels, kernel_size, stride, padding, bn_zero_init, bn_statistics_group_size, relu, name): """Initialize ConvBNRelu module. Args: out_channels (int): Number of output channels, i.e. number of convolution filters. kernel_size (int): Size of convolution kernel. stride (int): Convolution stride. padding (int): Convolution padding. bn_zero_init (bool): Zero-initialize batch normalization scale. bn_statistics_group_size (int): Group size for aggregating batch normalization statistics. relu (bool): Apply ReLU activation. name (str): Module name. """ super().__init__() self.name = name self.instance = 0 # Initialize convolution self.conv = lbann.modules.Convolution2dModule(out_channels, kernel_size, stride=stride, padding=padding, bias=False, name=self.name + '_conv') # Initialize batch normalization bn_scale_init = 0.0 if bn_zero_init else 1.0 bn_scale = lbann.Weights( initializer=lbann.ConstantInitializer(value=bn_scale_init), name=self.name + '_bn_scale') bn_bias = lbann.Weights( initializer=lbann.ConstantInitializer(value=0.0), name=self.name + '_bn_bias') self.bn_weights = [bn_scale, bn_bias] self.bn_statistics_group_size = bn_statistics_group_size # Initialize ReLU self.relu = relu
def __init__(self, size, bias=True, weights=[], activation=None, name=None, data_layout='data_parallel'): """Initialize fully-connected module. Args: size (int): Size of output tensor. activation (type): Layer class for activation function. bias (bool): Whether to apply bias after linearity. weights (`Weights` or iterator of `Weights`): Weights in fully-connected layer. There are at most two: the matrix and the bias. If weights are not provided, the matrix will be initialized with He normal initialization and the bias with zeros. name (str): Default name is in the form 'fcmodule<index>'. data_layout (str): Data layout. """ super().__init__() FullyConnectedModule.global_count += 1 self.instance = 0 self.size = size self.bias = bias self.name = (name if name else 'fcmodule{0}'.format(FullyConnectedModule.global_count)) self.data_layout = data_layout # Initialize weights # Note: If weights are not provided, matrix weights are # initialized with He normal scheme and bias weights are # initialized with zeros. self.weights = list(make_iterable(weights)) if len(self.weights) > 2: raise ValueError('`FullyConnectedModule` has ' 'at most two weights, ' 'but got {0}'.format(len(self.weights))) if len(self.weights) == 0: self.weights.append( lbann.Weights(initializer=lbann.HeNormalInitializer(), name=self.name+'_matrix')) if len(self.weights) == 1: self.weights.append( lbann.Weights(initializer=lbann.ConstantInitializer(value=0.0), name=self.name+'_bias')) # Initialize activation layer self.activation = None if activation: if isinstance(activation, type): self.activation = activation else: self.activation = type(activation) if not issubclass(self.activation, lbann.Layer): raise ValueError('activation must be a layer')
def forward_discriminator1(self, img): ''' Discriminator 1 ''' bn_wts = [ lbann.Weights(initializer=lbann.ConstantInitializer(value=1.0)), lbann.Weights(initializer=lbann.ConstantInitializer(value=0.0)) ] for count, lyr in enumerate(self.d1_conv): if count == 0: x = lbann.LeakyRelu(lyr(img), negative_slope=0.2) else: x = lbann.LeakyRelu(lyr(x), negative_slope=0.2) #### without convbrlu # if count==0: x = lbann.LeakyRelu(lbann.BatchNormalization(lyr(img),weights=bn_wts,statistics_group_size=-1),negative_slope=0.2) # else: x = lbann.LeakyRelu(lbann.BatchNormalization(lyr(x),weights=bn_wts,statistics_group_size=-1),negative_slope=0.2) dims = 32768 #dims=25088 ## for padding=1 y = self.d1_fc(lbann.Reshape(x, dims=str(dims))) return y
def __init__(self, embed_dim, num_heads, name=None): super().__init__() MultiheadAttention.global_count += 1 self.instance = 0 assert embed_dim % num_heads == 0, 'embed_dim must be divisible by num_heads' self.embed_dim = embed_dim self.num_heads = num_heads self.head_dim = embed_dim // num_heads # Module name self.name = name if not self.name: self.name = f'multiheadattention{MultiheadAttention.global_count}' # Weights for fully-connected layers self.query_weights = [ lbann.Weights(initializer=lbann.GlorotNormalInitializer(), name=f'{self.name}_query_matrix'), lbann.Weights(initializer=lbann.ConstantInitializer(value=0), name=f'{self.name}_query_bias'), ] self.key_weights = [ lbann.Weights(initializer=lbann.GlorotNormalInitializer(), name=f'{self.name}_key_matrix'), lbann.Weights(initializer=lbann.ConstantInitializer(value=0), name=f'{self.name}_key_bias'), ] self.value_weights = [ lbann.Weights(initializer=lbann.GlorotNormalInitializer(), name=f'{self.name}_value_matrix'), lbann.Weights(initializer=lbann.ConstantInitializer(value=0), name=f'{self.name}_value_bias'), ] self.output_weights = [ lbann.Weights(initializer=lbann.GlorotNormalInitializer(), name=f'{self.name}_output_matrix'), lbann.Weights(initializer=lbann.ConstantInitializer(value=0), name=f'{self.name}_output_bias'), ]
def __init__( self, embed_dim=512, num_heads=8, feedforward_dim=2048, dropout=0.1, name=None, ): TransformerDecoderLayer.global_count += 1 self.instance = 0 self.embed_dim = embed_dim self.feedforward_dim = feedforward_dim self.dropout_prob = dropout # Module name self.name = name if not self.name: self.name = f'transformerdecoderlayer{TransformerDecoderLayer.global_count}' # Layer modules self.attention1 = lbann.modules.transformer.MultiheadAttention( embed_dim, num_heads, name=f'{self.name}_attention1') self.attention2 = lbann.modules.transformer.MultiheadAttention( embed_dim, num_heads, name=f'{self.name}_attention2') # Weights for fully-connected layers self.fc1_weights = [ lbann.Weights(initializer=lbann.HeNormalInitializer(), name=f'{self.name}_fc1_matrix'), lbann.Weights(initializer=lbann.ConstantInitializer(value=0), name=f'{self.name}_fc1_bias'), ] self.fc2_weights = [ lbann.Weights(initializer=lbann.GlorotNormalInitializer(), name=f'{self.name}_fc2_matrix'), lbann.Weights(initializer=lbann.ConstantInitializer(value=0), name=f'{self.name}_fc2_bias'), ]
def __init__(self, neuron_dims=[1000, 1000, 1000], activation=lbann.Relu, keep_prob=0.95, name=None): self.instance = 0 self.name = (name if name else 'combo{0}'.format(Combo.global_count)) #shared weights for drug 1 and 2 tracks shared_w = [] for i in range(len(neuron_dims)): shared_w.append( lbann.Weights(initializer=lbann.HeNormalInitializer(), name='drug_matrix' + str(i))) shared_w.append( lbann.Weights(initializer=lbann.ConstantInitializer(value=0.0), name='drug_bias' + str(i))) print("SHARED W ", type(shared_w)) self.geneT = TrackModule(neuron_dims, activation, keep_prob, name=self.name + 'gene_track') self.drug1T = TrackModule(neuron_dims, activation, keep_prob, shared_w, name=self.name + 'drug1_track') self.drug2T = TrackModule(neuron_dims, activation, keep_prob, shared_w, name=self.name + 'drug2_track') self.concatT = TrackModule(neuron_dims, activation, keep_prob, name=self.name + 'concat_track')
def __init__(self, input_channels, output_channels, num_nodes, bias=True, activation=lbann.Relu, name=None, parallel_strategy={}): """Initialize GCN layer Args: input_channels (int): The size of the input node features output_channels (int): The output size of the node features num_nodes (int): Number of vertices in the graph bias (bool): Whether to apply biases after weights transform activation (type): Activation leyer for the node features. If None, then no activation is applied. (default: lbann.Relu) name (str): Default name of the layer is GCN_{number} parallel_strategy (dict): Data partitioning scheme. """ super().__init__() ## Add variables self.input_channel_size = input_channels self.output_channel_size = output_channels self.num_nodes = num_nodes self.parallel_strategy = parallel_strategy self.instance = 0 self.is_distconv = False if parallel_strategy: if list(parallel_strategy.values()[0]) > 0: self.is_distconv = True ## Add Name for the components for the layer GCNConv.global_count += 1 self.name = (name if name else 'GCN_{}'.format(GCNConv.global_count)) weights = [] ## Initialize weights for the matrix value = math.sqrt(6 / (input_channels + output_channels)) weights.append( lbann.Weights(initializer=lbann.UniformInitializer(min=-value, max=value), name=self.name + '_weights')) ## Initialize bias variables self.has_bias = bias if (self.has_bias): weights.append( lbann.Weights(initializer=lbann.ConstantInitializer(value=0.0), name=self.name + '_bias_weights')) self.activation = None if activation: if isinstance(activation, type): self.activation = activation else: self.activation = type(actvation) if not issubclass(self.activation, lbann.Layer): raise ValueError('activation must be a layer') # Distconv channelwise fully connected expects 3D tensors as input # and output. This check adds an extra dimention to enable # channel-wise data partitioning self.output_channels = self.output_channel_size if self.is_distconv: self.output_channels = [1, self.output_channel_size] self.nn = \ ChannelwiseFullyConnectedModule(self.output_channels, bias=self.has_bias, weights=weights, activation=self.activation, name=self.name+"_FC_layer", parallel_strategy=self.parallel_strategy)
def __init__(self, input_width, output_size, name=None, use_bn=False, bn_statistics_group_size=None): """Initialize CosmFlow. Args: input_width (int): Size of each spatial dimension of input data. output_size (int): Size of output tensor. name (str, optional): Module name (default: 'cosmoflow_module<index>'). use_bn (bool): Whether or not batch normalization layers are used. bn_statistics_group_size (int): The number of samples for each batch-normalization group. """ CosmoFlow.global_count += 1 self.instance = 0 self.name = (name if name else 'cosmoflow_module{0}'.format( CosmoFlow.global_count)) self.input_width = input_width self.use_bn = use_bn assert self.input_width in [128, 256, 512] self.cp_params = [ { "type": "conv", "out_channels": 16, "kernel_size": 3, "stride": 1 }, { "type": "pool" }, { "type": "conv", "out_channels": 32, "kernel_size": 3, "stride": 1 }, { "type": "pool" }, { "type": "conv", "out_channels": 64, "kernel_size": 3, "stride": 1 }, { "type": "pool" }, { "type": "conv", "out_channels": 128, "kernel_size": 3, "stride": 2 }, { "type": "pool" }, { "type": "conv", "out_channels": 256, "kernel_size": 3, "stride": 1 }, { "type": "pool" }, { "type": "conv", "out_channels": 256, "kernel_size": 3, "stride": 1 }, { "type": "conv", "out_channels": 256, "kernel_size": 3, "stride": 1 }, ] additional_pools = [] if self.input_width == 256: additional_pools = [6] elif self.input_width == 512: additional_pools = [6, 7] for i in additional_pools: conv_idx = list( np.cumsum([ 1 if x["type"] == "conv" else 0 for x in self.cp_params ])).index(i) self.cp_params.insert(conv_idx + 1, {"type": "pool"}) for p in self.cp_params: if p["type"] == "conv": p["padding"] = int((p["kernel_size"] - 1) / 2) # Create convolutional blocks activation = lbann.LeakyRelu for i, param in enumerate( filter(lambda x: x["type"] == "conv", self.cp_params)): conv_name = "conv" + str(i + 1) conv_weights = [ lbann.Weights(initializer=lbann.GlorotUniformInitializer()) ] param_actual = dict(param) param_actual.pop("type", None) conv = ConvBNRelu( **param_actual, conv_weights=conv_weights, use_bn=self.use_bn, bn_statistics_group_size=bn_statistics_group_size, bn_zero_init=False, name=self.name + "_" + conv_name, activation=lbann.LeakyRelu) setattr(self, conv_name, conv) # Create fully-connected layers fc_params = [ { "size": 2048 }, { "size": 256 }, { "size": output_size }, ] for i, param in enumerate(fc_params): fc_name = "fc" + str(i + 1) fc = lm.FullyConnectedModule( **param, activation=activation if i < len(fc_params) - 1 else None, name=self.name + "_" + fc_name, weights=[ lbann.Weights( initializer=lbann.GlorotUniformInitializer()), lbann.Weights(initializer=lbann.ConstantInitializer( value=0.1)) ], ) setattr(self, fc_name, fc)
def __init__(self, input_channels, output_channels, bias=True, activation=lbann.Relu, name=None, data_layout='data_parallel'): """Initialize Graph layer Args: input_channels (int): The size of the input node features output_channels (int): The output size of the node features bias (bool): Whether to apply biases after MatMul name (str): Default name of the layer is GCN_{number} data_layout (str): Data layout activation (type): Activation layer for the node features. If None, then no activation is applied. (default: lbann.Relu) """ super().__init__() ## Add variables self.input_channels = input_channels self.output_channels = output_channels self.data_layout = data_layout ## Add Name for the components for the layer GraphConv.global_count += 1 self.name = (name if name else 'Graph_{}'.format(GraphConv.global_count)) ## Initialize weights for the matrix value = math.sqrt(6 / (input_channels + output_channels)) self.mat_weights = lbann.Weights(initializer=lbann.UniformInitializer( min=-value, max=value), name=self.name + '_Weights') self.weights1 = lbann.WeightsLayer(dims=str_list( [input_channels, output_channels]), name=self.name + '_layer', weights=self.mat_weights) self.id_weights = lbann.Weights(initializer=lbann.UniformInitializer( min=-value, max=value), name=self.name + '_ID_Weights') self.weights2 = lbann.WeightsLayer(dims=str_list( [input_channels, output_channels]), name=self.name + '_ID_layer', weights=self.id_weights) ## Initialize bias variables self.has_bias = bias self.bias_weights = None self.bias = None if (self.has_bias): self.bias_weights = lbann.Weights( initializer=lbann.ConstantInitializer(value=0.0), name=self.name + '_bias_weights') self.bias = lbann.WeightsLayer(dims=str_list([1, output_channels]), weights=self.bias_weights, name=self.name + '_bias_layer') self.activation = None if activation: if isinstance(activation, type): self.activation = activation else: self.activation = type(actvation) if not issubclass(self.activation, lbann.Layer): raise ValueError('activation must be a layer')
def __init__(self, input_channels, output_channels, num_nodes, bias=True, activation=lbann.Relu, name=None): """Initialize Graph layer Args: input_channels (int): The size of the input node features output_channels (int): The output size of the node features num_nodes (int): Number of vertices in the graph bias (bool): Whether to apply biases after weights transform activation (type): Activation layer for the node features. If None, then no activation is applied. (default: lbann.Relu) name (str): Default name of the layer is Graph_{number} """ super().__init__() ## Add variables self.input_channel_size = input_channels self.output_channel_size = output_channels self.num_nodes = num_nodes ## Add Name for the components for the layer GraphConv.global_count += 1 self.name = (name if name else 'Graph_{}'.format(GraphConv.global_count)) ## Initialize weights for the matrix value = math.sqrt(6 / (input_channels + output_channels)) mat_weights = [] id_weights = [] mat_weights.append( lbann.Weights(initializer=lbann.UniformInitializer(min=-value, max=value), name=self.name + '_Weights')) id_weights.append( lbann.Weights(initializer=lbann.UniformInitializer(min=-value, max=value), name=self.name + '_ID_Weights')) ## Initialize bias variables self.has_bias = bias if (self.has_bias): mat_weights.append( lbann.Weights(initializer=lbann.ConstantInitializer(value=0.0), name=self.name + '_bias_weights')) self.activation = None if activation: if isinstance(activation, type): self.activation = activation else: self.activation = type(actvation) if not issubclass(self.activation, lbann.Layer): raise ValueError('activation must be a layer') self.id_nn = \ ChannelwiseFullyConnectedModule(self.output_channel_size, bias=False, weights=id_weights, activation=self.activation, name=self.name+"_ID_FC_layer") self.mat_nn = \ ChannelwiseFullyConnectedModule(self.output_channel_size, bias=self.has_bias, weights=mat_weights, activation=self.activation, name=self.name+"_Message_FC_layer")
def __init__(self, num_dims, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, weights=[], activation=None, name=None, transpose=False, parallel_strategy={}): """Initialize convolution module. Args: num_dims (int): Number of dimensions. out_channels (int): Number of output channels, i.e. number of filters. kernel_size (int): Size of convolution kernel. stride (int): Convolution stride. padding (int): Convolution padding. dilation (int): Convolution dilation. groups (int): Number of convolution groups. bias (bool): Whether to apply channel-wise bias after convolution. weights (`Weights` or iterator of `Weights`): Weights in convolution layer. There are at most two: the kernel and the bias. If weights are not provided, the kernel will be initialized with He normal initialization and the bias with zeros. name (str): Default name is in the form 'convmodule<index>'. transpose (bool): If true call deconvolution (or convolution transpose) parallel_strategy dict): Data partitioning scheme. """ super().__init__() ConvolutionModule.global_count += 1 self.instance = 0 self.num_dims = num_dims self.out_channels = out_channels self.kernel_size = kernel_size self.stride = stride self.padding = padding self.dilation = dilation self.groups = groups self.bias = bias self.weights = list(make_iterable(weights)) self.name = (name if name else 'convmodule{0}'.format( ConvolutionModule.global_count)) self.transpose = transpose self.parallel_strategy = parallel_strategy # Initialize weights # Note: If weights are not provided, kernel weights are # initialized with He normal scheme and bias weights are # initialized with zeros. self.weights = list(make_iterable(weights)) if len(self.weights) > 2: raise ValueError('`ConvolutionModule` has ' 'at most two weights, ' 'but got {0}'.format(len(self.weights))) if len(self.weights) == 0: self.weights.append( lbann.Weights(initializer=lbann.HeNormalInitializer(), name=self.name + '_kernel')) if len(self.weights) == 1: self.weights.append( lbann.Weights(initializer=lbann.ConstantInitializer(value=0.0), name=self.name + '_bias')) # Initialize activation layer self.activation = None if activation: if isinstance(activation, type): self.activation = activation else: self.activation = type(activation) if not issubclass(self.activation, lbann.Layer): raise ValueError('activation must be a layer')
'--work-dir', action='store', default=None, type=str, help='working directory', metavar='DIR') args = parser.parse_args() # ---------------------------------- # Embedding weights # ---------------------------------- encoder_embeddings_weights = lbann.Weights( initializer=lbann.NormalInitializer( mean=0, standard_deviation=1/args.latent_dim, ), name='embeddings', ) decoder_embeddings_weights = lbann.Weights( initializer=lbann.ConstantInitializer(value=0), name='decoder_embeddings', ) # ---------------------------------- # Construct layer graph # ---------------------------------- # Properties of graph and random walk num_graph_nodes = dataset.max_graph_node_id() + 1 walk_length = dataset.walk_context_length num_negative_samples = dataset.num_negative_samples input_size = dataset.sample_dims()[0] # Embedding vectors, including negative sampling # Note: Input is sequence of graph node IDs
def __init__(self, output_size, input_width, name=None): """Initialize CosmFlow. Args: output_size (int): Size of output tensor. input_width (int): Width of input tensor. name (str, optional): Module name (default: 'cosmoflow_module<index>'). """ CosmoFlow.global_count += 1 self.instance = 0 self.name = (name if name else 'cosmoflow_module{0}'.format( CosmoFlow.global_count)) self.input_width = input_width assert self.input_width in [128, 256, 512] self.layer_params = [ { "type": "conv", "out_channels": 16, "kernel_size": 3, "stride": 1 }, { "type": "pool" }, { "type": "conv", "out_channels": 32, "kernel_size": 3, "stride": 1 }, { "type": "pool" }, { "type": "conv", "out_channels": 64, "kernel_size": 3, "stride": 1 }, { "type": "pool" }, { "type": "conv", "out_channels": 128, "kernel_size": 3, "stride": 2 }, { "type": "pool" }, { "type": "conv", "out_channels": 256, "kernel_size": 3, "stride": 1 }, { "type": "pool" }, { "type": "conv", "out_channels": 256, "kernel_size": 3, "stride": 1 }, { "type": "conv", "out_channels": 256, "kernel_size": 3, "stride": 1 }, ] for p in self.layer_params: if p["type"] == "conv": p["padding"] = int((p["kernel_size"] - 1) / 2) additional_pools = [] if self.input_width == 256: additional_pools = [6] elif self.input_width == 512: additional_pools = [6, 7] for i in additional_pools: conv_idx = list( np.cumsum([ 1 if x["type"] == "conv" else 0 for x in self.layer_params ])).index(i) self.layer_params.insert(conv_idx + 1, {"type": "pool"}) width = self.input_width for p in self.layer_params: if p["type"] == "conv": output_width = int(width / p["stride"]) else: output_width = int(width / 2) p["width"] = output_width width = output_width assert width > 0 for i, param in enumerate( filter(lambda x: x["type"] == "conv", self.layer_params)): conv_name = "conv" + str(i + 1) conv_weights = [ Weights(initializer=lbann.GlorotUniformInitializer()) ] param_actual = dict(param) param_actual.pop("type", None) param_actual.pop("width", None) conv = lm.Convolution3dModule(**param_actual, activation=lbann.LeakyRelu, name=self.name + "_" + conv_name, bias=False, weights=conv_weights) setattr(self, conv_name, conv) # Create fully-connected layers fc_params = [ { "size": 2048 }, { "size": 256 }, { "size": output_size }, ] for i, param in enumerate(fc_params): fc_name = "fc" + str(i + 1) fc = lm.FullyConnectedModule( **param, activation=lbann.LeakyRelu if i < len(fc_params) - 1 else None, name=self.name + "_" + fc_name, weights=[ Weights(initializer=lbann.GlorotUniformInitializer()), Weights(initializer=lbann.ConstantInitializer(value=0.1)) ], ) setattr(self, fc_name, fc)