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)
