[WIP] federated knowledge distillation with heterogeneous model and molecule

federatedscope/gfkd/dataloader/dataloader_molecule.py

@@ -0,0 +1,87 @@
+from torch_geometric import transforms
+from torch_geometric.datasets import TUDataset, MoleculeNet, QM7b
+from rdkit import Chem
+from rdkit.Chem import AllChem
+from rdkit.Chem import Draw
+m = Chem.MolFromSmiles('c1ccccc1')
+m3d=Chem.AddHs(m)
+AllChem.EmbedMolecule(m3d, randomSeed=1)
+
+
+
+from federatedscope.core.auxiliaries.transform_builder import get_transform
+
+
+def load_heteromolecule_dataset(config=None):
+    r"""Convert dataset to Dataloader.
+    :returns:
+         data_local_dict
+    :rtype: Dict {
+                  'client_id': {
+                      'train': DataLoader(),
+                      'val': DataLoader(),
+                      'test': DataLoader()
+                               }
+                  }
+    """
+    splits = config.data.splits
+    path = config.data.root
+    name = config.data.type.upper()
+
+    # Transforms
+    transforms_funcs = get_transform(config, 'torch_geometric')
+
+    if name.startswith('heterogeneous molecule dataset'.upper()):
+        dataset = []
+
+        TUDdataset_names = ['BZR', 'ENZYMES', 'MUTAG']
+        MoleculeNet_names = ['ESOL', 'FreeSolv', 'BACE']
+        for dname in TUDdataset_names:
+            tmp_dataset = TUDataset(path, dname, **transforms_funcs)
+            dataset.append(tmp_dataset)
+        for dname in MoleculeNet_names:
+            tmp_dataset = MoleculeNet(path, dname, **transforms_funcs)
+
+            if dname in ['FreeSolv', 'BACE']:
+                for i in len(tmp_dataset):
+                    smiles = dataset[i].smiles
+                    mol = Chem.MolFromSmiles(smiles)
+                    mol = AllChem.AddHs(mol)
+                    res = AllChem.EmbedMolecule(mol, randomSeed=1)
+                    # will random generate conformer with seed equal to -1. else fixed random seed.
+                    if res == 0:
+                        try:
+                            AllChem.MMFFOptimizeMolecule(mol)# some conformer can not use MMFF optimize
+                        except:
+                            pass
+                        mol = AllChem.RemoveHs(mol)
+                        coordinates = mol.GetConformer().GetPositions()
+
+                    elif res == -1:
+                        mol_tmp = Chem.MolFromSmiles(smiles)
+                        AllChem.EmbedMolecule(mol_tmp, maxAttempts=5000, randomSeed=1)
+                        mol_tmp = AllChem.AddHs(mol_tmp, addCoords=True)
+                        try:
+                            AllChem.MMFFOptimizeMolecule(mol_tmp)# some conformer can not use MMFF optimize
+                        except:
+                            pass
+                        mol_tmp = AllChem.RemoveHs(mol_tmp)
+                        coordinates = mol_tmp.GetConformer().GetPositions()
+
+                    assert dataset[i].x.shape[0] == len(coordinates), "coordinates shape is not align with {}".format(smiles)
+                    tmp_dataset[i] = [tmp_dataset[i], coordinates]
+            dataset.append(tmp_dataset)
+        tmp_dataset = QM7b(path, dname, **transforms_funcs)
+        dataset.append(tmp_dataset)
+    else:
+        raise ValueError(f'No dataset named: {name}!')
+
+    client_num = min(len(dataset), config.federate.client_num
+                     ) if config.federate.client_num > 0 else len(dataset)
+    config.merge_from_list(['federate.client_num', client_num])
+
+    # get local dataset
+    data_dict = dict()
+    for client_idx in range(1, len(dataset) + 1):
+        data_dict[client_idx] = dataset[client_idx - 1]
+    return data_dict, config

federatedscope/gfkd/model/SMILES_model.py

@@ -0,0 +1,71 @@
+import torch
+import torch.nn
+import random
+import math
+
+class SMILESTransformer(torch.nn.Module):
+    def __init__(self, ntoken, ninp, nhead, nhid, nlayers, dropout=0.1):
+        super(SMILESTransformer, self).__init__()
+        from torch.nn import TransformerEncoder, TransformerEncoderLayer
+        self.model_type = 'Transformer'
+        self.ninp = ninp
+
+        self.pos_encoder = torch.nn.Embedding(100, ninp)
+        self.encoder = torch.nn.Embedding(ntoken, ninp)
+
+        self.layer_norm = torch.nn.LayerNorm([ninp])
+        self.output_layer_norm = torch.nn.LayerNorm([ntoken])
+        self.input_layer_norm = torch.nn.LayerNorm([ninp])
+
+        encoder_layers = TransformerEncoderLayer(d_model=ninp,
+                                                 nhead=nhead,
+                                                 dim_feedforward=nhid,
+                                                 dropout=dropout,
+                                                 activation='gelu')
+
+        self.transformer_encoder = TransformerEncoder(encoder_layers,
+                                                      nlayers,
+                                                      norm=self.layer_norm)
+
+        self.dropout = torch.nn.Dropout(dropout)
+
+        self.decoder = torch.nn.Linear(ninp, ntoken, bias=False)
+        self.decoder_bias = torch.nn.Parameter(torch.zeros(ntoken))
+        self.init_weights()
+
+
+    def init_weights(self):
+        initrange = 0.1
+        self.encoder.weight.data.normal_(mean=0.0, std=1.0)
+        self.decoder.weight.data.normal_(mean=0.0, std=1.0)
+        self.decoder_bias.data.zero_()
+
+        self.input_layer_norm.weight.data.fill_(1.0)
+        self.input_layer_norm.bias.data.zero_()
+        self.output_layer_norm.weight.data.fill_(1.0)
+        self.output_layer_norm.bias.data.zero_()
+        self.layer_norm.weight.data.fill_(1.0)
+        self.layer_norm.bias.data.zero_()
+
+
+    def forward(self, src, latent_out=False):
+        pos = torch.arange(0,100).long().to(src.device)
+
+        mol_token_emb = self.encoder(src)
+        pos_emb = self.pos_encoder(pos)
+        input_emb = pos_emb + mol_token_emb
+        input_emb = self.input_layer_norm(input_emb) 
+        input_emb = self.dropout(input_emb)
+        input_emb = input_emb.transpose(0, 1) 
+
+        attention_mask = torch.ones_like(src).to(src.device)
+        attention_mask = attention_mask.masked_fill(src!=1., 0.)
+        attention_mask = attention_mask.bool().to(src.device)
+
+        output = self.transformer_encoder(input_emb)
+
+        if latent_out:
+            return output
+        output = self.decoder(output) + self.decoder_bias 
+
+        return output