图神经网络应用：预测分子溶解度_溶解度模型和神经网络

ChatGPT狂飙160天，世界已经不是之前的样子。

新建了人工智能中文站https://ai.weoknow.com
每天给大家更新可用的国内可用chatGPT资源

---

在本篇博客中，我们将探讨如何使用化学结构来估计溶解度。这是一个回归问题，我们将提供SMILES字符串作为输入，并预测相应的溶解度。为了实现这个目标，我们将使用PyTorch Geometric构建图神经网络，并使用RDKit处理分子数据。

首先，我们需要安装PyTorch、PyTorch Geometric和RDKit。您可以使用以下命令通过pip进行安装：

pip install torch
pip install torch-geometric
pip install rdkit

数据集介绍

在接下来的内容中，我们将使用PyTorch Geometric的数据集库中的一个数据集。我们将使用MoleculeNet数据集中的ESOL数据集。ESOL数据集包含了1128种不同化学物质在水中的溶解度信息。该数据集被用于训练模型，通过直接从化学结构（以SMILES字符串编码）计算溶解度。由于溶解度通常是分子的特征，而不是特定构象的特征，因此这些结构不包括3D坐标。

我们的任务是研究溶质在溶剂中的溶解度。让我们先来看一下分子的SMILES表示。

加载数据集

import rdkit
from torch_geometric.datasets import MoleculeNet
 
# 加载ESOL数据集
data = MoleculeNet(root=".", name="ESOL")
data

获取数据集信息

print("Dataset type: ", type(data))
print("Dataset features: ", data.num_features)
print("Dataset target: ", data.num_classes)
print("Dataset length: ", data.len)
print("Dataset sample: ", data[0])
print("Sample  nodes: ", data[0].num_nodes)
print("Sample  edges: ", data[0].num_edges)

Dataset type:  <class 'torch_geometric.datasets.molecule_net.MoleculeNet'>
Dataset features:  9
Dataset target:  734
Dataset length:  <bound method InMemoryDataset.len of ESOL(1128)>
Dataset sample:  Data(x=[32, 9], edge_index=[2, 68], edge_attr=[68, 3], y=[1, 1], smiles='OCC3OC(OCC2OC(OC(C#N)c1ccccc1)C(O)C(O)C2O)C(O)C(O)C3O ')
Sample  nodes:  32
Sample  edges:  68

ESOL数据集包含了734个图，每个图代表一个化合物。图中的每个节点表示一个原子，边表示化学键。每个节点具有9个特征，目标维度也是734，表示每个图对应一个溶解度值。

# 查看图中节点的特征
data[0].x
 
# 查看边的信息（以稀疏COO格式表示）
# 形状为 [2, num_edges]
data[0].edge_index.t()
 
# 查看 data[0] 的目标值
data[0].y

在图级别上进行预测意味着我们要预测每个图所代表的化合物的溶解度值，而不是预测单个节点级别的属性。接下来，我们将把SMILES表示的分子转换为RDKit分子对象，并进行可视化。

data[0]["smiles"]

OCC3OC(OCC2OC(OC(C#N)c1ccccc1)C(O)C(O)C2O)C(O)C(O)C3O 

分子可视化

from rdkit import Chem
from rdkit.Chem.Draw import IPythonConsole
molecule = Chem.MolFromSmiles(data[0]["smiles"])
molecule

要使用RDKit从分子中提取特征，您可以利用RDKit提供的各种功能，例如提取边信息、原子属性等。在我们的情况下，由于数据集已经提供了所需的信息，因此使用RDKit提取特征更加简单。我们可以使用原子属性来计算节点特征。

图神经网络的实现

创建图神经网络的过程与创建卷积神经网络（Convolutional Neural Network）非常相似，只是我们添加了更多的层。图卷积网络（GCN）继承自torch.nn.Module。GCNConv是一个常用的图卷积层，它需要以下参数：

• in_channels：输入样本的大小。

• out_channels：输出样本的大小。

我们将使用三个图卷积层，这意味着我们可以获取三个邻居跳数的信息。为了进行图级别的预测，我们还会应用一个池化层来混合来自各个节点的数据。为了完成我们的任务，我们将使用PyTorch和PyTorch Geometric库。

以下是使用PyTorch和PyTorch Geometric实现图神经网络的示例代码：

import torch
from torch.nn import Linear
import torch.nn.functional as F 
from torch_geometric.nn import GCNConv, TopKPooling, global_mean_pool
from torch_geometric.nn import global_mean_pool as gap, global_max_pool as gmp
embedding_size = 64
 
class GCN(torch.nn.Module):
    def __init__(self):
        # Init parent
        super(GCN, self).__init__()
        torch.manual_seed(42)
 
        # GCN layers
        self.initial_conv = GCNConv(data.num_features, embedding_size)
        self.conv1 = GCNConv(embedding_size, embedding_size)
        self.conv2 = GCNConv(embedding_size, embedding_size)
        self.conv3 = GCNConv(embedding_size, embedding_size)
 
        # Output layer
        self.out = Linear(embedding_size*2, 1)
 
    def forward(self, x, edge_index, batch_index):
        # First Conv layer
        hidden = self.initial_conv(x, edge_index)
        hidden = F.tanh(hidden)
 
        # Other Conv layers
        hidden = self.conv1(hidden, edge_index)
        hidden = F.tanh(hidden)
        hidden = self.conv2(hidden, edge_index)
        hidden = F.tanh(hidden)
        hidden = self.conv3(hidden, edge_index)
        hidden = F.tanh(hidden)
 
        # Global Pooling (stack different aggregations)
        hidden = torch.cat([gmp(hidden, batch_index), 
                            gap(hidden, batch_index)], dim=1)
 
        # Apply a final (linear) classifier.
        out = self.out(hidden)
 
        return out, hidden
 
model = GCN()
print(model)
print("Number of parameters: ", sum(p.numel() for p in model.parameters()))

GCN(
  (initial_conv): GCNConv(9, 64)
  (conv1): GCNConv(64, 64)
  (conv2): GCNConv(64, 64)
  (conv3): GCNConv(64, 64)
  (out): Linear(in_features=128, out_features=1, bias=True)
)
Number of parameters:  13249

• 考虑到我们拥有大量的分子数据，我们选择使用64维的嵌入向量来表示它们，而不是将它们进行最小化处理。

• 随着我们逐渐增加网络层数，我们将能够更好地理解图的结构和特征。

• 为了解决回归问题，我们采用线性层作为最终的输出层。

• 尽管我们的样本数量只有大约1,000个，但我们致力于使用尽可能少的参数来实现高效的模型训练和预测。

训练神经网络

from torch_geometric.data import DataLoader
import warnings
warnings.filterwarnings("ignore")
 
# Root mean squared error
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.0007)  
 
# Use GPU for training
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
 
# Wrap data in a data loader
data_size = len(data)
NUM_GRAPHS_PER_BATCH = 64
loader = DataLoader(data[:int(data_size * 0.8)], 
                    batch_size=NUM_GRAPHS_PER_BATCH, shuffle=True)
test_loader = DataLoader(data[int(data_size * 0.8):], 
                         batch_size=NUM_GRAPHS_PER_BATCH, shuffle=True)
 
def train(data):
    # Enumerate over the data
    for batch in loader:
      # Use GPU
      batch.to(device)  
      # Reset gradients
      optimizer.zero_grad() 
      # Passing the node features and the connection info
      pred, embedding = model(batch.x.float(), batch.edge_index, batch.batch) 
      # Calculating the loss and gradients
      loss = loss_fn(pred, batch.y)     
      loss.backward()  
      # Update using the gradients
      optimizer.step()   
    return loss, embedding
 
print("Starting training...")
losses = []
for epoch in range(2000):
    loss, h = train(data)
    losses.append(loss)
    if epoch % 100 == 0:
      print(f"Epoch {epoch} | Train Loss {loss}")

Starting training...
Epoch 0 | Train Loss 3.377596378326416
Epoch 100 | Train Loss 0.9617947340011597
Epoch 200 | Train Loss 1.0771363973617554
Epoch 300 | Train Loss 0.6295697093009949
Epoch 400 | Train Loss 0.37517455220222473
Epoch 500 | Train Loss 0.465716689825058
Epoch 600 | Train Loss 0.5129485726356506
Epoch 700 | Train Loss 0.21677978336811066
Epoch 800 | Train Loss 0.33871856331825256
Epoch 900 | Train Loss 0.3640660345554352
Epoch 1000 | Train Loss 0.20501013100147247
Epoch 1100 | Train Loss 0.18023353815078735
Epoch 1200 | Train Loss 0.2812242805957794
Epoch 1300 | Train Loss 0.18207958340644836
Epoch 1400 | Train Loss 0.1321338415145874
Epoch 1500 | Train Loss 0.18665631115436554
Epoch 1600 | Train Loss 0.1817774772644043
Epoch 1700 | Train Loss 0.09456530958414078
Epoch 1800 | Train Loss 0.23615044355392456
Epoch 1900 | Train Loss 0.11381624639034271

可视化学习过程

import seaborn as sns
losses_float = [float(loss.cpu().detach().numpy()) for loss in losses] 
loss_indices = [i for i,l in enumerate(losses_float)] 
plt = sns.lineplot(loss_indices, losses_float)
plt

测试集上进行预测

import pandas as pd 
 
# Analyze the results for one batch
test_batch = next(iter(test_loader))
with torch.no_grad():
    test_batch.to(device)
    pred, embed = model(test_batch.x.float(), test_batch.edge_index, test_batch.batch) 
    df = pd.DataFrame()
    df["y_real"] = test_batch.y.tolist()
    df["y_pred"] = pred.tolist()
df["y_real"] = df["y_real"].apply(lambda row: row[0])
df["y_pred"] = df["y_pred"].apply(lambda row: row[0])
df

下面是可视化y_pred与y_original的代码：

plt = sns.scatterplot(data=df, x="y_real", y="y_pred")
plt.set(xlim=(-7, 2))
plt.set(ylim=(-7, 2))
plt

以上是我们应用第一个GNN的示例代码。

---

ChatGPT狂飙160天，世界已经不是之前的样子。

新建了人工智能中文站https://ai.weoknow.com
每天给大家更新可用的国内可用chatGPT资源