def __init__(self, in_size, out_size=20, hidden_size=128, neighbours=30, w_depth=3, p_depth=3, max_distance=20, distance_kernels=16): super().__init__() self.features = MLP(2 * in_size, hidden_size, hidden_size=hidden_size, depth=p_depth, batch_norm=False, normalization=spectral_norm) self.weight = MLP(2 * in_size, 1, hidden_size=hidden_size, depth=w_depth, batch_norm=False, normalization=spectral_norm) self.out = MLP(hidden_size * neighbours, out_size, hidden_size=[300, hidden_size, hidden_size], batch_norm=False, normalization=spectral_norm) self.rbf = (0, max_distance, distance_kernels)
def __init__(self, aa_size=21, in_size=5, hidden_size=100, p_depth=3, w_depth=2, neighbours=10): super(Baseline, self).__init__() self.in_size = in_size self.probability = MLP(2 * in_size, aa_size, hidden_size=hidden_size, depth=p_depth, batch_norm=False) self.weight = MLP(2 * in_size, 1, hidden_size=hidden_size, depth=w_depth, batch_norm=False) self.out = MLP(aa_size * neighbours, aa_size, hidden_size=[300, hidden_size, hidden_size], batch_norm=False)
def __init__(self): super().__init__() self.process = MLP(28 * 28, 64, 128, depth=3, normalization=spectral_norm, batch_norm=False) self.predict = MLP(64, 10, 128, depth=3, normalization=spectral_norm) self.condition = MLP(1, 128, 128, depth=3, normalization=spectral_norm) self.scale = nn.Linear(128, 64) self.bias = nn.Linear(128, 64)
def __init__(self, depth=4): super(ConvEnergy, self).__init__() self.preprocess = nn.Conv2d(1, 32, 1) self.blocks = nn.ModuleList([ spectral_norm(nn.Conv2d(32, 32, 3, padding=1)) for idx in range(depth) ]) self.bn = nn.ModuleList([ nn.GroupNorm(8, 32) for idx in range(depth) ]) self.postprocess = nn.Conv2d(32, 128, 1) self.condition = MLP(10, 128, depth=3, batch_norm=False, normalization=spectral_norm) self.combine = MLP(128, 1, hidden_size=64, depth=3, batch_norm=False, normalization=spectral_norm)
def __init__(self, size, distance_size, attention_size=128, heads=128, hidden_size=128, mlp_depth=3, activation=func.relu_, batch_norm=False, dropout=0.1, pre_norm=True, normalization=lambda x: x, connected=attention_connected): super(ReverseDistanceBlock, self).__init__() self.batch_norm = batch_norm self.attention = ts.NeighbourDotMultiHeadAttention( distance_size + 2, size, attention_size, query_size=size, heads=heads, normalization=normalization ) self.local = MLP( size, size, hidden_size=hidden_size, depth=mlp_depth, activation=activation, batch_norm=False, normalization=normalization ) self.activation = activation self.dropout = lambda x: x if dropout is not None: self.dropout = nn.Dropout(dropout) self.bn = lambda x: x self.local_bn = lambda x: x if self.batch_norm: self.bn = nn.LayerNorm(size) self.local_bn = nn.LayerNorm(size)
def __init__(self): super().__init__() self.input = MLP(28 * 28 + 100, 28 * 28, hidden_size=128, depth=3, batch_norm=False)
def __init__(self, size, distance_size, attention_size=128, heads=128, hidden_size=128, mlp_depth=3, activation=func.relu_, batch_norm=False, dropout=0.1, pre_norm=True, normalization=lambda x: x, connected=attention_connected): super(StructuredTransformerEncoderBlock, self).__init__() self.pre_norm = pre_norm self.batch_norm = batch_norm self.attention = connected(size, distance_size, attention_size, heads, normalization) self.local = MLP( size, size, hidden_size=hidden_size, depth=mlp_depth, activation=activation, batch_norm=False, normalization=normalization ) self.activation = activation self.dropout = lambda x: x if dropout is not None: self.dropout = nn.Dropout(dropout) self.bn = lambda x: x self.local_bn = lambda x: x if self.batch_norm: self.bn = nn.LayerNorm(size) self.local_bn = nn.LayerNorm(size)
def __init__(self, in_size, size, distance_size, sequence_size=20, attention_size=128, heads=128, hidden_size=128, mlp_depth=3, depth=3, max_distance=20, distance_kernels=16, neighbours=15, activation=func.relu_, batch_norm=True, conditional=False, angles=False, dropout=0.1, connected=attention_connected, normalization=spectral_norm): super().__init__() distance_size = distance_size + distance_kernels - 1 self.encoder = StructuredTransformerEncoder( size, size, distance_size, attention_size=attention_size, heads=heads, hidden_size=hidden_size, depth=depth, mlp_depth=mlp_depth, activation=activation, batch_norm=batch_norm, dropout=dropout, normalization=normalization, connected=connected ) self.angles = angles self.lookup = PositionLookup(fragment_size=10) self.conditional = conditional self.neighbours = neighbours self.activation = activation self.rbf = (0, max_distance, distance_kernels) self.preprocess = LocalFeatures(6, size) self.postprocess = LocalFeatures(size, size) self.result = nn.Linear(2 * size, 1) self.angle_result = MLP(6, 1, hidden_size=32, depth=3, batch_norm=False, normalization=spectral_norm)
def __init__(self, depth=4): super(Connected, self).__init__() self.block = MLP(28 * 28, 1, depth=depth, batch_norm=False, normalization=spectral_norm)
def __init__(self, in_size, hidden_size, depth=4, max_dilation=4): super(ResidualLocalInteraction, self).__init__() self.global_mlp = MLP(2 * in_size, in_size, hidden_size=in_size) self.blocks = nn.ModuleList([ ResidualBlock(in_size, hidden_size, dilation=2 ** (idx % 4)) for idx in range(depth) ])
def __init__(self, single, size=5, latents=16): super(Encoder, self).__init__() self.size = size self.single = single self.weight = spectral_norm(nn.Linear(32, 1)) self.combine = MLP(32, 32, 64, depth=3, batch_norm=False, normalization=spectral_norm) self.mean = spectral_norm(nn.Linear(32, latents)) self.logvar = spectral_norm(nn.Linear(32, latents))
def __init__(self, single, size=5, latents=16): super(Encoder, self).__init__() self.size = size self.single = single self.weight = nn.Linear(32, 1) self.combine = MLP(32, 32, 64, depth=3) self.mean = nn.Linear(32, latents) self.logvar = nn.Linear(32, latents)
def __init__(self): super(Energy, self).__init__() self.input = MLP(28 * 28 + 10, 28 * 28, hidden_size=128, depth=4, batch_norm=False, activation=func.elu)
def __init__(self, sample=True): super(Energy, self).__init__() self.sample = sample self.input = SingleEncoder() self.condition = Encoder(self.input) self.input_process = spectral_norm(nn.Linear(32, 64)) self.postprocess = spectral_norm(nn.Linear(16, 64)) self.combine = MLP(128, 1, hidden_size=64, depth=4, batch_norm=False, normalization=spectral_norm)
def __init__(self): super(Critic, self).__init__() self.input = MLP(32 * 32 * 3 + 10, 32 * 32 * 3, hidden_size=512, depth=4, batch_norm=False, activation=swish) self.embed = nn.Linear(10, 10)
def __init__(self): super(Energy, self).__init__() self.input = MLP(28 * 28, 128, hidden_size=128, depth=3, batch_norm=False, normalization=spectral_norm) self.condition = MLP(10, 128, depth=3, batch_norm=False, normalization=spectral_norm) self.combine = MLP(128, 1, hidden_size=64, depth=3, batch_norm=False, normalization=spectral_norm)
def __init__(self): super().__init__() self.energy = nn.Parameter( torch.randn(1) ) #MLP(10, 1, 64, depth=3, normalization=spectral_norm, batch_norm=False) self.predict = MLP(28 * 28, 10, 128, depth=3, normalization=spectral_norm, batch_norm=False)
def __init__(self, encoder, size=5): super(Discriminator, self).__init__() self.size = size self.encoder = encoder self.input = encoder.single self.verdict = MLP(28 * 28 + 16, 1, hidden_size=128, depth=4, batch_norm=False, activation=func.leaky_relu)
def __init__(self): super(Energy, self).__init__() self.input = MLP(32 * 32 * 3 + 10, 1, hidden_size=1000, depth=5, batch_norm=False, activation=swish) self.embed = nn.Linear(10, 10) self.mean = nn.Linear(10, 1, bias=False) self.logvar = nn.Linear(10, 1, bias=False) with torch.no_grad(): self.mean.weight.zero_() self.logvar.weight.zero_()
def __init__(self, size, distance_size, sequence_size, attention_size=128, heads=128, hidden_size=128, mlp_depth=3, activation=func.relu_, batch_norm=False, adaptive=None, dropout=0.1, pre_norm=True, normalization=lambda x: x): super(StructuredTransformerDecoderBlock, self).__init__() self.batch_norm = batch_norm self.pre_norm = pre_norm self.adaptive = adaptive self.attention = ts.NeighbourDotMultiHeadAttention( size + distance_size + sequence_size, size, attention_size, query_size=size, heads=heads, normalization=normalization) self.local = MLP(size, size, hidden_size=hidden_size, depth=mlp_depth, activation=activation, batch_norm=False, normalization=normalization) self.activation = activation self.dropout = lambda x: x if dropout is not None: self.dropout = nn.Dropout(dropout, inplace=True) self.bn = lambda x: x self.local_bn = lambda x: x if self.batch_norm: if self.adaptive is not None: self.bn = AdaptiveLayerNorm(size, adaptive) self.local_bn = AdaptiveLayerNorm(size, adaptive) else: self.bn = nn.LayerNorm(size) self.local_bn = nn.LayerNorm(size)
def __init__(self, input_net, sequence_size=32, hidden_size=64, latent_size=64, depth=3, activation=func.elu_): super(TorsionConditionalPrior, self).__init__() mlp_hidden = 2 * latent_size self.condition = input_net self.activation = activation self.postprocess = MLP(hidden_size, mlp_hidden, hidden_size=mlp_hidden, depth=depth - 1, activation=activation) self.mu = nn.Linear(mlp_hidden, latent_size) self.sigma = nn.Linear(mlp_hidden, latent_size)
def __init__(self, angle_net, input_net, sequence_size=32, latent_size=64, depth=3, activation=func.elu_, return_torsion=False): super(TorsionConditionalDecoder, self).__init__() mlp_hidden = 2 * latent_size torsion_factor = 4 - 2 * return_torsion self.sequence_size = sequence_size self.angle = angle_net self.condition = input_net self.postprocess = MLP(mlp_hidden, sequence_size * torsion_factor, hidden_size=mlp_hidden, depth=depth - 1, activation=activation)
def __init__(self): super().__init__() self.mapping = MLP(512, 512, hidden_size=128, depth=8, batch_norm=False) start_features = 512 self.blocks = nn.ModuleList([ StyleGANBlock(start_features, 128, start_features, size=(4, 4), activation=func.relu_) ] + [ StyleGANBlock(128, 128, start_features, activation=func.relu_) for idx in range(6) ]) self.postprocess = nn.Conv2d(128, 1, 3, padding=1)
def __init__(self, distance_size=128, hidden_size=128, angles=512, fragment_size=5, attention_size=128, heads=8, depth=3, mlp_depth=3, dropout=0.1, activation=func.relu_, batch_norm=False, pre_norm=True, normalization=lambda x: x, connected=attention_connected): super(DistanceDiscriminator, self).__init__() self.angle_lookup = nn.Linear(3, hidden_size) self.rbf = (0, 20, 64) self.distance_lookup = MLP(64, hidden_size, hidden_size=hidden_size, depth=3, batch_norm=False) self.angle_result = nn.Linear(hidden_size, 1) self.distance_result = nn.Linear(hidden_size, 1) self.preproc = LocalFeatures(hidden_size, hidden_size) self.transformer = DistanceTransformer( hidden_size, hidden_size, hidden_size, hidden_size=hidden_size, attention_size=attention_size, heads=8, depth=depth, mlp_depth=mlp_depth, dropout=dropout, activation=activation, batch_norm=batch_norm, pre_norm=pre_norm, normalization=normalization, connected=connected )
def __init__(self, angle_net, input_net, sequence_size=32, latent_size=64, depth=3, activation=func.elu_, return_torsion=False): super(TorsionDrawDecoder, self).__init__(angle_net, input_net, sequence_size=sequence_size, latent_size=latent_size, depth=depth, activation=activation, return_torsion=return_torsion) mlp_hidden = 2 * latent_size torsion_factor = 4 - 2 * return_torsion self.predict_angle = MLP(mlp_hidden, sequence_size * torsion_factor, hidden_size=mlp_hidden, depth=depth - 1, activation=activation)
def __init__(self, in_size, out_size, hidden_size, query_size=None, attention_size=64, heads=8, mlp_depth=3, dropout=0.1, activation=func.elu_): super(Transformer, self).__init__(in_size, out_size, hidden_size, query_size=query_size, attention_size=attention_size, heads=heads, dropout=dropout, activation=activation) self.local_block = MLP(in_size, hidden_size, depth=mlp_depth, batch_norm=False, activation=activation)
def __init__(self, size, n_heads=8, hidden_size=128, attention_size=128, value_size=128, depth=2, dropout=0.1): super().__init__() self.attention = SequenceMultiHeadAttention( size, size, attention_size=attention_size, hidden_size=value_size, heads=n_heads) self.ff = MLP(size, size, hidden_size=hidden_size, depth=2, batch_norm=False) self.rezero = ReZero(size) self.dropout_1 = nn.Dropout(dropout) self.dropout_2 = nn.Dropout(dropout)
logits=torch.zeros(15, dtype=torch.float), outputs=None ) def forward(self, inputs, **kwargs): result = self.blocks(inputs) result = func.adaptive_avg_pool2d(result, 1).view(result.size(0), -1) result = self.postprocess(result) return result mode = sys.argv[1] index = sys.argv[2] training = ... if mode == "bdpi": policy = Policy() value = MLP(4, 2, batch_norm=False, depth=3) agent = CategoricalPolicy(SharedModule(policy)) env = CoinRun() training = BDPITraining( policy, value, agent, env, network_name=f"awr-test/bdpi-coinrun-{index}", discount=0.99, clones=4, critic_updates=4, gradient_updates=20, batch_size=128, buffer_size=100_000, device="cuda:0", verbose=True )
def __init__(self, latents=32): super(SingleEncoder, self).__init__() self.block = MLP(28 * 28, latents, hidden_size=64, depth=4, batch_norm=False, normalization=spectral_norm)
def __init__(self, z=32): super().__init__() self.z = z self.generate = MLP(3 * 28 * 28 + z, 3 * 28 * 28, depth=4)