Skip to content

Commit

Permalink
escn so2 exports
Browse files Browse the repository at this point in the history
  • Loading branch information
@rayg1234
rayg1234 committed Aug 29, 2024
1 parent e4a426a commit 672e1ff
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 36 deletions.
21 changes: 11 additions & 10 deletions src/fairchem/core/models/escn/escn.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -777,8 +777,8 @@ def forward(
x_target._rotate(SO3_edge_rot, self.lmax_list, self.mmax_list)

# Compute messages
x_source = self.so2_block_source(x_source, x_edge)
x_target = self.so2_block_target(x_target, x_edge)
x_source.embedding = self.so2_block_source(x_source.embedding, x_edge)
x_target.embedding = self.so2_block_target(x_target.embedding, x_edge)

# Add together the source and target results
x_target.embedding = x_source.embedding + x_target.embedding
Expand Down Expand Up @@ -859,41 +859,42 @@ def forward(
num_edges = len(x_edge)

# Reshape the spherical harmonics based on m (order)
x._m_primary(self.mappingReduced)
x = torch.einsum("nac,ba->nbc", x, self.mappingReduced.to_m)

# Compute m=0 coefficients separately since they only have real values (no imaginary)

# Compute edge scalar features for m=0
x_edge_0 = self.act(self.fc1_dist0(x_edge))

x_0 = x.embedding[:, 0 : self.mappingReduced.m_size[0]].contiguous()
x_0 = x[:, 0 : self.mappingReduced.m_size[0]].contiguous()
x_0 = x_0.view(num_edges, -1)

x_0 = self.fc1_m0(x_0)
x_0 = x_0 * x_edge_0
x_0 = self.fc2_m0(x_0)
x_0 = x_0.view(num_edges, -1, x.num_channels)
x_0 = x_0.view(num_edges, -1, self.sphere_channels)

# Update the m=0 coefficients
x.embedding[:, 0 : self.mappingReduced.m_size[0]] = x_0
x[:, 0 : self.mappingReduced.m_size[0]] = x_0

# Compute the values for the m > 0 coefficients
offset = self.mappingReduced.m_size[0]
for m in range(1, max(self.mmax_list) + 1):
# Get the m order coefficients
x_m = x.embedding[
x_m = x[
:, offset : offset + 2 * self.mappingReduced.m_size[m]
].contiguous()
x_m = x_m.view(num_edges, 2, -1)
# Perform SO(2) convolution
x_m = self.so2_conv[m - 1](x_m, x_edge)
x_m = x_m.view(num_edges, -1, x.num_channels)
x.embedding[:, offset : offset + 2 * self.mappingReduced.m_size[m]] = x_m
x_m = x_m.view(num_edges, -1, self.sphere_channels)
x[:, offset : offset + 2 * self.mappingReduced.m_size[m]] = x_m

offset = offset + 2 * self.mappingReduced.m_size[m]

# Reshape the spherical harmonics based on l (degree)
x._l_primary(self.mappingReduced)
# x._l_primary(self.mappingReduced)
x = torch.einsum("nac,ab->nbc", x, self.mappingReduced.to_m)

return x

Expand Down
59 changes: 33 additions & 26 deletions tests/core/models/test_escn_compiles.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -110,18 +110,22 @@ def test_escn_compiles(self):
# torch._dynamo.config.suppress_errors = True

# os.environ["TORCH_LOGS"] = "+dynamo,recompiles"
# torch._logging.set_logs(dynamo = logging.INFO)
torch._logging.set_logs(dynamo = logging.INFO)
# os.environ["TORCHDYNAMO_VERBOSE"] = "1"
# os.environ["TORCHDYNAMO_REPRO_AFTER"]="dynamo"
# torch._dynamo.config.verbose = True
compiled_model = torch.compile(model, dynamic=True)
torch._dynamo.config.optimize_ddp = False
# torch._dynamo.explain(model)(data)
# assert False
# torch._dynamo.reset()
# explain_output = torch._dynamo.explain(model)(data)
# print(explain_output)

output = compiled_model(data)
expected_energy, expected_forces = expected_energy_forces()
assert torch.allclose(output["energy"], expected_energy)
assert torch.allclose(output["forces"].mean(0), expected_forces)
# expected_energy, expected_forces = expected_energy_forces()
# assert torch.allclose(output["energy"], expected_energy)
# assert torch.allclose(output["forces"].mean(0), expected_forces)

def test_rotation_invariance(self) -> None:
random.seed(1)
Expand Down Expand Up @@ -150,25 +154,28 @@ def test_rotation_invariance(self) -> None:
decimal=5,
)

# def test_escn_so2_conv_compiles(self) -> None:
# torch._dynamo.config.assume_static_by_default = False
# torch._dynamo.config.automatic_dynamic_shapes = True
# inp1_dim0 = Dim("inp1_dim0")
# inp1_dim1 = None
# inp1_dim2 = None
# inp2_dim0 = inp1_dim0
# inp2_dim1 = Dim("inp2_dim1")

# dynamic_shapes1 = {
# "x_emb": {0: inp1_dim0, 1: inp1_dim1, 2: inp1_dim2},
# "x_edge": {0: inp2_dim0, 1: inp2_dim1},
# }

# lmax, mmax = 4, 2
# mappingReduced = CoefficientMapping([lmax], [mmax])
# edge_channels = 128
# args=(torch.rand(680, 19, 128), torch.rand(680, edge_channels))

# so2 = SO2Block(sphere_channels=128, hidden_channels=128, edge_channels=edge_channels, lmax_list=[lmax], mmax_list=[mmax], act=torch.nn.SiLU())
# prog = export(so2, args=args, dynamic_shapes=dynamic_shapes1)
# export_out = prog.module()(*args)
def test_escn_so2_conv_exports(self) -> None:
torch._dynamo.config.assume_static_by_default = False
torch._dynamo.config.automatic_dynamic_shapes = True
inp1_dim0 = Dim("inp1_dim0")
inp1_dim1 = None
inp1_dim2 = None
inp2_dim0 = inp1_dim0
inp2_dim1 = None

dynamic_shapes1 = {
"x": {0: inp1_dim0, 1: inp1_dim1, 2: inp1_dim2},
"x_edge": {0: inp2_dim0, 1: inp2_dim1},
}

lmax, mmax = 4, 2
mappingReduced = CoefficientMapping([lmax], [mmax])
shpere_channels = 128
edge_channels = 128
args=(torch.rand(680, 19, shpere_channels), torch.rand(680, edge_channels))

so2 = SO2Block(sphere_channels=shpere_channels, hidden_channels=128, edge_channels=edge_channels, lmax_list=[lmax], mmax_list=[mmax], act=torch.nn.SiLU(), mappingReduced=mappingReduced)
prog = export(so2, args=args, dynamic_shapes=dynamic_shapes1)
export_out = prog.module()(*args)
regular_out = so2(*args)
assert torch.allclose(export_out, regular_out)

0 comments on commit 672e1ff

Please sign in to comment.