Paper再现：MD+AI自动编码机探测蛋白变构(三)：训练自编码模型

三、训练自编码模型

文章中是用的自动编码机是很简答的模型，如下图。由于这一部分是比较简单的，有一点神经网络基础的都很快可以完成，所以就写的简单一些哈。

我们也来模仿这个简单的自动编码机模型。使用pytorch来建模。

3.1 导入相关模块

#导入相关模块
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
from torch.optim import lr_scheduler
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
torch.set_default_tensor_type(torch.DoubleTensor)
device

输出为：device(type='cuda', index=0) 因为，我这里使用的是带一个GPU的机子。

构建自动编码机模型

（很简单，大家可以尝试其他的模型）

class Aotuencoder_for_md(nn.Module):
    
    def __init__(self):
        super(Aotuencoder_for_md, self).__init__()
        
        self.encoder_layer1 = nn.Linear(1000, 800)
        self.encoder_layer2 = nn.Linear(800, 500)
        self.encoder_layer3 = nn.Linear(500, 200)
        
        self.decoder_layer1 = nn.Linear(200, 500)
        self.decoder_layer2 = nn.Linear(500, 800)
        self.decoder_layer3 = nn.Linear(800, 1000)
    
    def forward(self, x):
        y = F.relu(self.encoder_layer1(x))
        y = F.relu(self.encoder_layer2(y))
        y = F.relu(self.encoder_layer3(y))
        
        y = F.relu(self.decoder_layer1(y))
        y = F.relu(self.decoder_layer2(y))
        y = F.relu(self.decoder_layer3(y))
        return y

3.2 实例化模型

model = Aotuencoder_for_md()
model = model.to(device)

3.3 加载数据

主要是加载我们在第二部分数据预处理中处理好的数据，主要是：opo_distance_dataset.npy和holo_distance_dataset.npy两个文件，里面做了一些处理，例如删除了全为0的残基对（自己对自己的残疾对），同时由于模型是全连接神经网络，因此，对数据集进行了最大最小归一化。

import numpy as np
import torch.utils.data as Data
from sklearn.model_selection import train_test_split

opo = np.load('./dataset/dataset/opo_distance_dataset.npy')
print('opo 初始残基对数量：', len(opo))
mask = np.all(np.isnan(opo), axis=1) | np.all(opo == 0, axis=1)
opo = opo[~mask]
opo = opo  / opo.max(axis=0)
print('opo 删除self-self残基对后数量：',len(opo))

#加载测试集数据
holo = np.load('./dataset/dataset/holo_distance_dataset.npy')
print('holo 初始残基对数量：', len(holo))
mask = np.all(np.isnan(holo), axis=1) | np.all(holo == 0, axis=1)
holo = holo[~mask]
holo = holo / holo.max(axis=0)
print('holo 删除self-self残基对后数量：',len(holo))

#划分测试集和训练集
opo_train, opo_test = train_test_split(opo, test_size=0.3, random_state=42)

#数据加载器
opo_dataset_train = Data.DataLoader(torch.from_numpy(opo_train), batch_size=100,
                                     shuffle=True, num_workers=64, pin_memory=True)
opo_dataset_test = Data.DataLoader(torch.from_numpy(opo_test), batch_size=100,
                                    shuffle=True, num_workers=64, pin_memory=True)
holo_dataset = Data.DataLoader(torch.from_numpy(holo), batch_size=100,
                                    shuffle=True, num_workers=64, pin_memory=True)

输出为：

样本的数据量还是挺大的，opo有8万多条残基对，holo有9万多。

3.4 设置损失函数和优化器

这里我们使用的学习率阶梯下降，尽可能的按照原文的设置去做。

from torch import optim
optimizer = optim.Adam(model.parameters(), lr=0.01)
#学习率下降
scheduler = lr_scheduler.StepLR(optimizer,step_size=100*len(opo_dataset_train),gamma = 0.9)

#损失函数
loss_func = nn.MSELoss()

3.5 日志函数

由于训练过程很长，为了观察，写一个打日志文件

def run_log(t, file='run_lof.txt'):
    f = open (file, 'a+')
    f.write(t)
    time_ = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) 
    f.write(time_)
    f.write('\n')

3.6 训练模型

早期停止部分有点偷懒了，建议大家根据需求修改

#训练集和测试集的损失
history = {'train_loss':[], 'test_loss':[]}
#训练opo1模型
for epoch in range(10000):
    train_loss=0
    model.train()
    for i, x in enumerate(opo_dataset_train):
        x = x.to(device)
        y_pred = model(x)
        loss = loss_func(y_pred, x)
        train_loss = train_loss + loss.item()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        scheduler.step()        
#     print('{} optim: {}'.format(epoch, optimizer.param_groups[0]['lr'])) #查看学习率
    train_loss = round(train_loss / len(opo_dataset_train), 6)
    history['train_loss'].append(train_loss)
    
    test_loss=0
    model.eval()
    for _, x in enumerate(opo_dataset_test):
        x = x.to(device)
        y_pred = model(x)
        loss = loss_func(y_pred, x)
        test_loss = test_loss + loss.item()
    test_loss = round(test_loss / len(opo_dataset_test), 6)
    history['test_loss'].append(test_loss)
    
    #如果关掉了窗口，则print函数不会接续输出训练的状态
#     print('训练次数：{} optim: {}'.format(epoch, optimizer.param_groups[0]['lr']))
        
    if epoch % 10 == 0:
        log = '训练次数：{}， opo1训练集损失为：{}， 测试集损失为：{}'.format(epoch, train_loss, test_loss)
        print(log)
        run_log(log)
        print('训练次数：{} optim: {}'.format(epoch, optimizer.param_groups[0]['lr']))
        torch.save(model, './model.pth') 
        
    #早期停止训练，patience为1000 epochs
    if epoch == 0:
        k = 0
        min_loss = test_loss
    else:
        if test_loss >= min_loss:
            k = k + 1
        else:
            k = 0
            min_loss = test_loss
    
    if k > 1000:
        print('EarlyStop training.')
        log = 'EarlyStop training.'
        run_log(log)
        break

输出为：

训练完成后，工作目录下的model.pth就是训练好的模型。
训练过程还是很耗时间的，文章中直接是设置到了3万步迭代。虽然只是1玩不，还带1000的早期停止，即使我是用的是V100的卡，也训练了两天。。。

接下来就是DIO的生成与分析了。