<jupyter_start><jupyter_text>In this notebook I'm going to build MLP as predictor on physiochemical properties and fingerprints. The MLP structure comes from `SmipleMLP` object in TAPE project.<jupyter_code>import numpy as np
import pandas as pd
import torch
from torch import nn
import torch.nn.functional as F
from torch.nn.utils.weight_norm import weight_norm
from tqdm import tqdm, trange
from sklearn.metrics import mean_absolute_error
import matplotlib.pyplot as plt
%matplotlib inline
%load_ext autoreload
%autoreload 2
import ivpk
class SimpleRegHead(nn.Module):
    def __init__(self,
                 in_dim: int,
                 hid_dim: int, 
                 dropout: float = 0., 
                 intersept: float = None):
        super().__init__()
        self.main = nn.Sequential(
            nn.Linear(in_dim, hid_dim), 
            nn.GELU(), 
            nn.Dropout(dropout, inplace=False),
            nn.Linear(hid_dim, 1)
        )
        # set the bias if needed
        if intersept:
            with torch.no_grad():
                self.main[3].bias.fill_(intersept)

    def forward(self, x):
        return self.main(x)
device = "cuda" if torch.cuda.is_available() else "cpu"<jupyter_output><empty_output><jupyter_text>Use VDss as example target and physiochemical properties as example input.<jupyter_code>target = 'VDss'<jupyter_output><empty_output><jupyter_text>Dataset and data loader. Be careful that the dimensionality for y should not be (N, ) but (N, 1) here. Otherwise the loss function will not calculate properly due to broadcasting rules.<jupyter_code>train_dl, val_dl = ivpk.data.get_dataloaders(target, ("train", "val"), smiles_func=None)<jupyter_output><empty_output><jupyter_text>Dummy run training.<jupyter_code>reg = SimpleRegHead(11, 512).to(device)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(reg.parameters(), lr=0.01)
def loss_batch(model, loss_fn, xb, yb, opt=None):
    loss = loss_fn(model(xb), yb)

    if opt is not None:
        loss.backward()
        opt.step()
        opt.zero_grad()

    return loss.item(), len(xb)

def fit(num_epochs, model, loss_fn, optimizer, train_dl, val_dl):
    for epoch in trange(num_epochs):
        model.train()
        losses, nums = zip(
            *[loss_batch(model, loss_fn, Xb, yb, optimizer) for Xb, yb in train_dl]
        )
        train_loss = np.sum(np.multiply(losses, nums)) / np.sum(nums)
        model.eval()
        with torch.no_grad():
            losses, nums = zip(
                *[loss_batch(model, loss_fn, Xb, yb) for Xb, yb in val_dl]
            )
        val_loss = np.sum(np.multiply(losses, nums)) / np.sum(nums)
        yield (epoch, train_loss, val_loss)
history = pd.DataFrame(
    fit(500, reg, loss_fn, optimizer, train_dl, val_dl), 
    columns=['epoch', 'train_loss', 'val_loss']
)
plt.plot(history.epoch, history.train_loss)
plt.plot(history.epoch, history.val_loss)
plt.show()
x_val, y_val = val_dl.dataset.tensors
plt.plot(y_val, reg(x_val).detach().numpy(), ".")
x_train, y_train = train_dl.dataset.tensors
plt.plot(y_train, reg(x_train).detach().numpy(), ".")
print(
    "MAE on val", 
    mean_absolute_error(y_val, reg(x_val).detach().numpy()), 
    "\nMAE on train", 
    mean_absolute_error(y_train, reg(x_train).detach().numpy())
)<jupyter_output>MAE on val 1.1275502 
MAE on train 1.1229836<jupyter_text>Now we have a basic idea of MLP performance on physiochemical properties. Let's try adding morgan fingerprint.<jupyter_code>train_dl, val_dl = ivpk.data.get_dataloaders(target, ("train", "val"), fpType="morgan", fpSize=256)
reg = SimpleRegHead(256+11, 512).to(device)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(reg.parameters(), lr=0.01)
history = pd.DataFrame(
    fit(500, reg, loss_fn, optimizer, train_dl, val_dl), 
    columns=['epoch', 'train_loss', 'val_loss']
)
plt.plot(history.epoch, history.train_loss)
plt.plot(history.epoch, history.val_loss)
plt.show()
x_val, y_val = val_dl.dataset.tensors
plt.plot(y_val, reg(x_val).detach().numpy(), ".")
x_train, y_train = train_dl.dataset.tensors
plt.plot(y_train, reg(x_train).detach().numpy(), ".")
print(
    "MAE on val", 
    mean_absolute_error(y_val, reg(x_val).detach().numpy()), 
    "\nMAE on train", 
    mean_absolute_error(y_train, reg(x_train).detach().numpy())
)<jupyter_output>MAE on val 1.1538448 
MAE on train 0.5062531<jupyter_text>With morgan256 we can overfit the training set, which is a good sign for the model capacity. Regularization can be explored later, e.g. dropout.  Now let's see if CL can also be properly predicted.<jupyter_code>target = "CL"
train_dl, val_dl = ivpk.data.get_dataloaders(target, ("train", "val"), smiles_func=None)
#train_dl, val_dl = ivpk.data.get_dataloaders(target, ("train", "val"), fpType="morgan", fpSize=256)
reg = SimpleRegHead(11, 512).to(device)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(reg.parameters(), lr=0.01)
history = pd.DataFrame(
    fit(500, reg, loss_fn, optimizer, train_dl, val_dl), 
    columns=['epoch', 'train_loss', 'val_loss']
)
plt.plot(history.epoch, history.train_loss)
plt.plot(history.epoch, history.val_loss)
plt.show()
x_val, y_val = val_dl.dataset.tensors
plt.plot(y_val, reg(x_val).detach().numpy(), ".")
x_train, y_train = train_dl.dataset.tensors
plt.plot(y_train, reg(x_train).detach().numpy(), ".")
print(
    "MAE on val", 
    mean_absolute_error(y_val, reg(x_val).detach().numpy()), 
    "\nMAE on train", 
    mean_absolute_error(y_train, reg(x_train).detach().numpy())
)
train_dl, val_dl = ivpk.data.get_dataloaders(target, ("train", "val"), fpType="morgan", fpSize=256)
reg = SimpleRegHead(11+256, 512).to(device)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(reg.parameters(), lr=0.01)
history = pd.DataFrame(
    fit(500, reg, loss_fn, optimizer, train_dl, val_dl), 
    columns=['epoch', 'train_loss', 'val_loss']
)
plt.plot(history.epoch, history.train_loss)
plt.plot(history.epoch, history.val_loss)
plt.show()
x_val, y_val = val_dl.dataset.tensors
plt.plot(y_val, reg(x_val).detach().numpy(), ".")
x_train, y_train = train_dl.dataset.tensors
plt.plot(y_train, reg(x_train).detach().numpy(), ".")
print(
    "MAE on val", 
    mean_absolute_error(y_val, reg(x_val).detach().numpy()), 
    "\nMAE on train", 
    mean_absolute_error(y_train, reg(x_train).detach().numpy())
)<jupyter_output>MAE on val 1.5796382 
MAE on train 0.47783652
{'fraction_non_alphanumeric': 0.14121897580898524, 'fraction_numerical': 0.018221803330191643, 'mean_word_length': 4.122284794851167, 'pattern_counts': {'":': 0, '<': 22, '<?xml version=': 0, '>': 22, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '440968', 'license': 'MIT', 'n_tokens_mistral': 2457, 'n_tokens_neox': 2337, 'n_words': 511, 'path': 'MLP_head.ipynb', 'repo_name': 'yanshouyu/human-iv-pk', 'size': '6366'}