生物信息实验bioinfo_lab

前言

        本实验来自于山东大学计算机科学与技术学院生物信息学2023秋季学期课程，仅供参考，勿作他用；实验内有缺陷，读者应该设计出更好的实验。本博客最终解释权归本人所有。

实验

0.广播机制

        这一实验主要学习numpy和pytorch的广播机制，广播机制对矩阵乘法非常重要，务必掌握得十分熟练，对之后的实验很有帮助。关键代码如下：

distance=pt.sqrt(pt.sum((coord[None,:,:]-coord[:,None,:])**2,dim=-1))

1.手写数字分类（手搓）

        在MNIST上训练一个神经网络，这个全连接神经网络需要自己实现。关键代码如下：

for x, y in tqdm(train_data):
    hi = x @ self.w1 + self.b1                    # [batch_size,784]@[784,256]=[64,256]
    ho = pt.where(hi > 0, hi, pt.zeros_like(hi))                        # ReLU [64,256]

    fi = ho @ self.w2 + self.b2                             # [64,256]@[256,10]=[64,10]
    fo = pt.softmax(fi, dim=1)                                                # [64,10]
    # ∵ dfo/dfi = e^fi(Σe^j-e^fi)/(Σe^j)^2 = fo(1-fo)

    loss = -y*pt.log(fo)+(y-1.0)*pt.log(1.0-fo)   
    running_loss.append(loss)                            # cross entropy [64,10]

    # 又∵ dloss/dfo = -y/fo + (1-y)/(1-fo)
    # ∴ dloss/dfi = -y(1-fo) + (1-y)fo = -y+fo
    grad_fi = -y+fo                                                            # [64,10]
    grad_w2 = ho.reshape((-1, 256, 1)) @ grad_fi.reshape((-1, 1, 10))          # 256×10
    grad_b2 = grad_fi  # [64,10]

    grad_ho = grad_fi @ self.w2.t()                           # [64,10]@[10,256]=[64,256]
    # 原来≤0的导数为0，原来＞0的导数为1
    grad_hi = grad_ho * pt.where(hi > 0, pt.ones_like(hi), pt.zeros_like(hi))  # [64,256]
    grad_w1 = x.reshape((-1, 784, 1)) @ grad_hi.reshape((-1, 1, 256))          # [784,256]
    grad_b1 = grad_ho                                                          # [64,256]
    # 使用梯度的平均值来更新参数可以更好地代表整个训练集的梯度方向，减小参数更新的方差，并减少训练过程中的震荡现象，从而提高模型的训练效果和稳定性
    self.w2 -= self.lr*pt.mean(grad_w2,axis=0)
    self.b2 -= self.lr*pt.mean(grad_b2,axis=0)
    self.w1 -= self.lr*pt.mean(grad_w1,axis=0)
    self.b1 -= self.lr*pt.mean(grad_b1,axis=0)

2. 手写数字分类（调接口，可视化）

        这次实验和上次的不同主要体现在调用现成接口实现，并对梯度进行可视化，关键代码如下：

class NeuralNetwork(nn.Module):
    def __init__(self,inodes,onodes,lr,hnodes=256) -> None:
        super(NeuralNetwork,self).__init__()
        self.lr=lr
        self.ful1=nn.Linear(inodes,hnodes)
        self.ful2=nn.Linear(hnodes,onodes)
        self.module=nn.Sequential(self.ful1,nn.ReLU(),self.ful2)
        self.loss_fun=nn.CrossEntropyLoss() #会自动加上softmax，先softmax，然后log
        self.optim=pt.optim.SGD(self.module.parameters(),lr=self.lr)

    def forward(self,x):
        return self.module(x)

3.CATH超家族分类（构建序列蛋白质数据的dataset）

        这次实验需要写一个提取数据的dataset以及批处理函数collate_fn，之后实验4、5都会用到，关键代码如下：

class ProteinDataset(Dataset):
    def __init__(self, dataset, mapping=None):
        super(ProteinDataset, self).__init__()
        if isinstance(dataset,tuple): # raw data
            self.seq=dataset[0]
            self.idx=dataset[1]
            self.lab=dataset[2]
            self.map=np.arange(len(self.lab),dtype=np.int64) # 恒等映射
            assert len(self.seq)==self.idx[-1]
            assert len(self.lab)==len(self.idx)-1
        else: # structured data
            self.seq=dataset.seq
            self.idx=dataset.idx
            self.lab=dataset.lab
            self.map=mapping
            assert self.map is not None
            assert np.max(self.map)<len(self.lab)
    # self.map旨在维护一个data子集的映射，数据仍然是全部数据   
    def __getitem__(self, idx):
        idx_=self.map[idx]
        seq=self.seq[self.idx[idx_]:self.idx[idx_+1]]
        seq=np.concatenate([[21],seq,[22]]) # add start and end token
        lab=self.lab[idx_]
        return deepcopy(seq),deepcopy(lab)

    def __len__(self):
        return len(self.map) #子集的大小是map的大小


def collate_fn(batch, bucketsize=batchsize*6):
    seq = [i[0] for i in batch] + [[0] * (bucketsize - 1)]
    seq = np.concatenate(seq)
    seq = seq[:len(seq)//bucketsize*bucketsize] #截取bucketsize的整数倍
    # ptr是每个样本序列（包括【21】【22】）长度的前缀和 
    ptr = [0] + [len(i[0]) for i in batch]
    ptr = np.append(np.cumsum(ptr), [len(seq)]) 
    # lab是batch中每个样本的标签，注意这里是i[1]，与数据结构中的getitem对应
    lab = [i[1] for i in batch]
    lab = np.array(lab)
    return seq.astype(np.int16), ptr.astype(np.int32), lab.astype(np.int16)

4.CATH超家族分类（构建CNN模型）

        这次实验需要手搓一个针对一维数据的CNN模型，原理见我之前的博客：Conv1d一维卷积（unfold），关键代码如下：

class EmbedBlock(nn.Module):
    def __init__(self, num_to_pad=400) -> None:
        super(EmbedBlock, self).__init__()
        self.pad = num_to_pad

    def forward(self, seq_ptr):
        seq = seq_ptr[0]
        ptr = seq_ptr[1]
        seq_ = []
        for i in range(len(ptr) - 2):
            s = seq[ptr[i] + 1:ptr[i + 1] - 1]  # 去掉开始和结束的21,22
            zero_pad = self.pad - len(s)
            if zero_pad % 2 == 0:
                s = np.insert(s, 0, [0] * (zero_pad // 2))
                s = np.append(s, [0] * (zero_pad // 2))
            else:
                s = np.insert(s, 0, [0] * (zero_pad // 2))
                s = np.append(s, [0] * (zero_pad // 2 + 1))
            seq_.append(s)
        seq_ = np.array(seq_)

        return from_numpy(seq_).float().unsqueeze(1).cuda()


class Conv1d(nn.Module):
    def __init__(self, in_channels:int, out_channels:int, kernel_size:int=3, 
                 stride:int=1, padding:int=0, bias:bool=True) -> None:
        super(Conv1d, self).__init__()
        self.stride = stride
        self.pad = padding
        self.ker_size = kernel_size        
        self.out_ch = out_channels
        self.kernel = nn.Parameter(nn.init.xavier_normal_(pt.randn(out_channels, in_channels, kernel_size)))
        self.bias = None
        if bias:
            self.bias = nn.Parameter(pt.zeros(out_channels))

    def forward(self, x):    
        # x：(batch_size, in_channel, in_len)
        # out_len = (x.shape[2] + 2*self.pad - self.ker_size)//self.stride + 1
        x = nn.functional.pad(x, (self.pad, self.pad), "constant", 0)
        
        # (batch_size, in_ch, 1, in_len) -> (batch_size, in_ch*ker_size, out_len)
        x_unfold = nn.functional.unfold(x[:,:,None,:], 
                                        kernel_size=(1,self.ker_size), stride=(1,self.stride))
        out = self.kernel.view(self.out_ch, -1).matmul(x_unfold)
                 
        if self.bias is not None:
            out += self.bias[None,:,None]
 
        return out

        PS：其实对蛋白质的idx序列进行卷积是不合理的，因为idx只表示不同种类的氨基酸，不同种类的氨基酸的差异并不是通过idx大小差异来衡量的，是离散型数据，而卷积针对的是连续型数据。

        正确思路是：使用pytorch提供的nn.Embedding接口，那么对于一个有n个氨基酸的蛋白质，就可以生成n个嵌入维度为x的embeddings，再对n这个维度进行加和或者求平均，得到1个embedding，这个embedding可以做一维卷积。这种思路请读者自己实践。

5. CATH超家族分类（构建MHA模型）

        这次实验需要手搓一个MHA模型，原理见我之前的博客：Multi-Head Self-Attention以及Self-Attention(einsum)，关键代码如下：

class SelfAttention(nn.Module):
    def __init__(self, d_m, d_k) -> None:
        super(SelfAttention, self).__init__()
        self.d_m = d_m
        self.d_k = d_k
        self.Wq = nn.Linear(in_features=self.d_m, out_features=self.d_k)
        self.Wk = nn.Linear(in_features=self.d_m, out_features=self.d_k)
        self.Wv = nn.Linear(in_features=self.d_m, out_features=self.d_k)

    def forward(self, x):
        q,k,v = self.Wq(x), self.Wk(x),self.Wv(x)
        score = pt.einsum('nci, ncj -> nc',q,k)
        score /= pt.sqrt(pt.tensor(self.d_k))
        score = nn.functional.softmax(score, dim=-1)
        out = v * score[:,:,None]
        return out
    

class MultiHeadSelfAttention(nn.Module):
    def __init__(self, embed_dim:int, num_heads=1) -> None:
        super(MultiHeadSelfAttention, self).__init__()
        assert embed_dim % num_heads == 0, 'embed_dim must be divided by num_heads'
        self.heads = nn.ModuleList([SelfAttention(d_m=embed_dim, d_k=embed_dim//num_heads) 
                                    for _ in range(num_heads)])
        self.Wo = nn.Linear(in_features=embed_dim, out_features=embed_dim)
    
    def forward(self, x):
        Z = pt.cat([head(x) for head in self.heads], dim=-1)
        if len(self.heads) > 1:
            return self.Wo(Z)
        else:
            return Z

6.CATH超家族分类（构建空间结构蛋白质数据的dataset）

        这次实验用到的数据新增了氨基酸α碳原子的空间坐标，以及氨基酸之间的连接关系，所以数据集的构建与实验3有些区别，关键代码如下：

class ProteinDataset(Dataset):
    def __init__(self, dataset, mapping=None):
        super(ProteinDataset, self).__init__()
        if isinstance(dataset, tuple): # raw data
            self.seq = pt.tensor(dataset[0], dtype=pt.int32)
            self.node_pos = pt.tensor(dataset[1], dtype=pt.float32)
            self.node_idx = pt.tensor(dataset[2])
            self.edge_nho = pt.tensor(dataset[3], dtype=pt.int32)
            self.edge_idx = pt.tensor(dataset[4])
            self.lab = pt.tensor(dataset[5])
            self.map = pt.arange(len(self.lab), dtype=pt.int32) # 恒等映射 
            assert len(self.seq) == self.node_idx[-1]
            assert len(self.lab) == len(self.node_idx) - 1
            assert len(self.lab) == len(self.edge_idx) - 1
        else: # structured data
            self.seq = dataset.seq
            self.node_pos = dataset.node_pos
            self.node_idx = dataset.node_idx
            self.edge_nho = dataset.edge_nho
            self.edge_idx = dataset.edge_idx
            self.lab = dataset.lab
            self.map = mapping
            assert self.map is not None
            assert pt.max(self.map) < len(self.lab)
    # self.map旨在维护一个data子集的映射，数据仍然是全部数据   
    
    def get(self, idx):
        idx_ = self.map[idx]
        seq = self.seq[self.node_idx[idx_] : self.node_idx[idx_+1]]
        len_seq = len(seq)
        seq_ = pt.zeros(len_seq, 21, dtype=pt.float32)
        seq_[pt.arange(len_seq) , seq[:]-1] = 1.0
        # shape:(len_seq,5,3) 0:N, 1:α, 2:C, 3:O, 4:β 
        node_pos = self.node_pos[self.node_idx[idx_] : self.node_idx[idx_+1]] 
        # N-α-β这个夹角反映了氨基酸的空间结构，在化学键确定的前提下，N和β之间的距离就能反应角度 
        seq_[:, 20] = pt.sqrt(pt.sum((node_pos[:,0] - node_pos[:,4])**2, dim=1))
        
        # 连接关系(氢键) 
        edge_nho = pt.stack((self.edge_nho[0][self.edge_idx[idx_] : self.edge_idx[idx_+1]],
                             self.edge_nho[1][self.edge_idx[idx_] : self.edge_idx[idx_+1]]), dim=0)
        # 连接关系(肽键)
        edge_tai = pt.stack((pt.arange(0, len_seq-1), pt.arange(1, len_seq)), dim=0)
        # 连接关系
        edge_idx = pt.cat((edge_nho, edge_tai), dim=1)
        # 边长 
        nho_attr = pt.empty(edge_nho.shape[1])
        # 氢键：氨基的氢和羧基的氧之间吸引产生，用氨基的氮坐标代替氢坐标 
        nho_attr[:] = pt.sqrt(pt.sum((node_pos[edge_nho[0][:], 0] - node_pos[edge_nho[1][:], 3])**2, dim=1))
        tai_attr = pt.empty(edge_tai.shape[1])
        # 羧基碳接氨基氮
        tai_attr[:] = pt.sqrt(pt.sum((node_pos[edge_tai[0][:], 2] - node_pos[edge_tai[1][:], 0])**2, dim=1))
        edge_attr = pt.cat((nho_attr, tai_attr))[:, None]
        
        lab = self.lab[idx_]
        one_hot = pt.zeros(6630, dtype=pt.float32)
        one_hot[lab[0]] = 1.0
        graph = Data(x=seq_, edge_index=edge_idx, edge_attr=edge_attr, y=one_hot, pos=node_pos[:,1])

        return graph
        

    def len(self):  
        return len(self.map) #子集的大小是map的大小

7.CATH超家族分类（构建MPNN模型）

        这次实验需要构建一个消息传递神经网络，我实现的是GCN，GCN的论文理解见我之前的博客：GCN图卷积神经网络论文阅读，以及搭建：GCN图卷积神经网络搭建，关键代码如下：

class HeteroGCNConv(nn.Module):
    def __init__(self, in_channels, out_channels) -> None:
        super(HeteroGCNConv, self).__init__()
        self.lin = nn.Linear(in_channels, out_channels, bias=False)
        self.conv1 = GCNConv(out_channels)
        self.conv2 = GCNConv(out_channels)

    def forward(self, x, edge_index, edge_attr):
        # edge_index:[2, N-1+num_nho] edge_attr:[N-1+num_nho, 1]
        N = x.size(0)
        tai_idx = pt.stack((edge_index[0, :N-1], edge_index[1, :N-1]), dim=0)
        tai_idx = pt.cat((tai_idx, tai_idx.flip(0)), dim=1)
        nho_idx = pt.stack((edge_index[0, N-1:], edge_index[1, N-1:]), dim=0)
        nho_idx = pt.cat((nho_idx, nho_idx.flip(0)), dim=1)
        tai_attr = edge_attr[:N-1]
        tai_attr = pt.cat((tai_attr, tai_attr), dim=0)
        nho_attr = edge_attr[N-1:]
        nho_attr = pt.cat((nho_attr, nho_attr), dim=0)
        
        x = self.lin(x)
        x = F.relu(x)
        x = self.conv1(x, tai_idx, tai_attr)
        x = F.relu(x)
        x = self.conv2(x, nho_idx, nho_attr)

        return x


class ProteinGCN(nn.Module):
    def __init__(self, in_channels:int, batch_size:int, out_channels=6630) -> None:
        super(ProteinGCN, self).__init__() 
        self.batch_size =  batch_size
        self.out_channels = out_channels
        self.conv1 = HeteroGCNConv(in_channels, 128)
        self.conv2 = HeteroGCNConv(128, 256)
        self.conv3 = HeteroGCNConv(256, 512)
        self.conv4 = HeteroGCNConv(512, 1024)
        self.JK = gnn.JumpingKnowledge('cat', [128, 256, 512, 1024])
        self.lin = nn.Linear(1920, out_channels)
    
    def forward(self, data):
        x, edge_index, edge_attr, batch = data.x, data.edge_index, data.edge_attr, data.batch
        x1 = self.conv1(x, edge_index, edge_attr)
        x2 = self.conv2(x1, edge_index, edge_attr)
        x3 = self.conv3(x2, edge_index, edge_attr)
        x4 = self.conv4(x3, edge_index, edge_attr)
        x = self.JK([x1, x2, x3, x4])
        x = gnn.global_mean_pool(x, batch)
        x = self.lin(x)

        return x

项目

        这个项目是做一个蛋白质的查询数据库以查找同源的蛋白质，本人当时没有做，这里提供一个思路。

• 训练数据应该是三个字母表示一种氨基酸，测试数据应该是一个字母表示一种氨基酸；
• 建议把训练数据处理为lab7的格式，这样就可以套用之前的模板了。
• 训练的目的是生成能独特地表示每一种蛋白质的embeddings，然后就可以用embedding对蛋白质进行检索，这里推荐一个好用的检索工具：脸书的faiss。

结语

        这里更新不方便，完整代码以GitHub为准：bioinfo。

        日后惹出事来，不要说是在这里找到的。