FraGAT：一种针对分子特征学习的图神经网络 源码解读分析 2021SC@SDUSC

2021SC@SDUSC
FraGAT：一种针对分子图特征提取的GNN，对于我们的DTA任务的分子表示学习部分有很大的启发，下面将学习分析它的代码与训练过程。

任务类型：线性回归

输入：药物分子SMILES
预测输出：分子属性值

FraGAT网络架构如下，它是分子经过断裂处理后分为四个部分经过特征提取最后拼接而获得最终的分子表示，我们的DTA任务的分子特征提取部分可以借鉴这一点，所以我们的重点将放在如何提取concat之后，进入FCN之前的特征向量。

代码分析：

从论文作者的训练实例开始分析（ESOL_sample.py）：

数据集截图，形式为一个SMILES加一个分子属性分数值

首先代码的开头是一个参数的字典ParamList

ParamList = {
   
   
    'ExpName': 'ESOL',
    'MainMetric': 'RMSE',
    'DataPath': './data/ESOL_SMILESValue.txt',
    'RootPath': './Experiments/',
    'CUDA_VISIBLE_DEVICES': '1',
    'TaskNum': 1,
    'ClassNum': 1,
    'Augmentation': False,
    'Weight': True,


    'ValidRate': 4000,
    'PrintRate': 20,
    'Frag': True,
    'output_size': 1,
    'atom_feature_size': 39,
    'bond_feature_size': 10,
    'Feature': 'AttentiveFP',

    'ValidBalance': False,
    'TestBalance': False,
    'MaxEpoch': 800,
    'SplitRate': [0.8,0.1],
    'Splitter': 'Random',

    'UpdateRate': 1,
    'LowerThanMaxLimit': 50,
    'DecreasingLimit': 30,

    'FP_size': 150,
    'atom_layers':3, #原子层是3
    'mol_layers':2, #分子层是2
    'DNNLayers':[512],
    'BatchSize':200,#批次大小为200
    'drop_rate':0.2,
    'lr':2.5,
    'WeightDecay':5,

    'SplitValidSeed': 38,
    'SplitTestSeed': 8,
    'TorchSeed': 2
}

接着是训练前的准备，将参数列表封装给Configs、创建保存结果的文件夹、检查cuda是否可用

os.environ['CUDA_VISIBLE_DEVICES'] = ParamList['CUDA_VISIBLE_DEVICES']

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print (torch.cuda.device_count())

#将参数列表装入config对象并赋值给opt
opt = Configs(ParamList)
print(opt.args['DataPath'])

#创建实验结果保存的文件夹
opt.add_args('SaveDir', opt.args['RootPath'] + opt.args['ExpName'] + '/')
print(opt.args['SaveDir'])
if not os.path.exists(opt.args['SaveDir']):
    os.mkdir(opt.args['SaveDir'])
model_dir = opt.args['SaveDir'] + 'model/'
if not os.path.exists(model_dir):
    os.mkdir(model_dir)

print("cuda available?"+torch.cuda.is_available()) #打印检查cuda是否可用

调用train_and_evaluate函数进行训练和评估，我们接着来看这个函数内部

ckpt, value = train_and_evaluate(opt) #训练与评估

train_and_evaluate函数完整代码，后面对其进行探究

def train_and_evaluate(opt):

    #指定seed(用于初始化神经网络参数)
    t.manual_seed(opt.args['TorchSeed'])
    saver = Saver(opt)

    #加载网络，优化器，以及批次
    net, optimizer, StartEpoch = saver.LoadModel()

    print(not net)


    #若不存在现有模型，则根据配置选择下面的初始模型架构
    if not net:

        #若参数feature的值是graph
        if opt.args['Feature'] == 'Graph':
            #39,恰好对应点的特征长度 该层是GCN层
            net = MolPredGraph(feature_size=39, GCN_layers=opt.args['GCNLayers'], DNN_layers=opt.args['DNNLayers'],
                               GCNLayertype=opt.args['GCNLayertype'])
            #若feature参数为FP的话
        elif opt.args['Feature'] == 'FP':

            net = MolPredFP(feature_size=opt.args['feature_size'], DNN_layers=opt.args['DNNLayers'])
            #main函数里调用的话应该是执行这个（相当于把opt里的参数赋值给网络的各个环节）
        elif opt.args['Feature'] == 'AttentiveFP':
            #测试main里Frag是等于true的（可能为false就不使用碎片化
            if opt.args['Frag'] == True:
                net = MolPredFragFPv8(
                    #原子特征 39
                    atom_feature_size=opt.args['atom_feature_size'],
                    #边特征 10
                    bond_feature_size=opt.args['bond_feature_size'],
                    #FP的长度 150（是提取的nodeEmbedding的大小吗？
                    FP_size=opt.args['FP_size'],
                    #原子层
                    atom_layers=opt.args['atom_layers'],
                    #分子层
                    mol_layers=opt.args['mol_layers'],
                    #DNN层（全连接层？
                    DNN_layers=opt.args['DNNLayers'],
                    #输出y的size（肯定只有一个呀
                    output_size=opt.args["output_size"],
                    drop_rate=opt.args['drop_rate'],
                    opt=opt
                )
            else:
                net = MolPredAttentiveFP(
                    atom_feature_size=opt.args['atom_feature_size'],
                    bond_feature_size=opt.args['bond_feature_size'],
                    FP_size=opt.args['FP_size'],
                    atom_layers=opt.args['atom_layers'],
                    mol_layers=opt.args['mol_layers'],
                    DNN_layers=opt.args['DNNLayers'],
                    output_size=opt.args["output_size"],
                    drop_rate=opt.args['drop_rate']
                )

    #将所有模型参数和缓存送入GPU
    net = net.cuda()

    #分子数据集创建 MolDatasetCreator会从opt中获取需要加载的数据路径并加载处理
    moldatasetcreator = MolDatasetCreator(opt)

    #数据集分割率
    #'SplitRate': [0.8,0.1],
    if len(opt.args['SplitRate']) == 1:
        (Trainset, Validset), weights = moldatasetcreator.CreateDatasets()