jkarolczak / ligand-classification Goto Github PK

import os

import torch
from datetime import datetime

def save_state_dict(
    model: torch.nn.Module,
    directory: str,
    epoch: int
) -> None:
    """
    Serialize model weigts.
    :param model: torch.nn.Module to save its weights
    :param directory: path to the directory to log experiment results
    :param epoch: int describing epoch
    """
    models_path = os.path.join(directory, 'models')
    os.makedirs(models_path, exist_ok=True)
    time = str(datetime.now()).replace(' ', '-')
    file_name = f'{time}-epoch-{epoch}.pt'
    file_path = os.path.join(models_path, file_name)
    torch.save(model.state_dict(), file_path)

import os

import numpy as np
import pandas as pd

import torch
from torch import nn
from torch import tensor
from torch.utils.data import Dataset

import MinkowskiEngine as ME

### datset and dataloader ###

class LigandDataset(Dataset):
    """A class to represent a ligands dataset."""
    def __init__(
        self, 
        annotations_file_path: str,
        labels_file_path: str = None
        ):
        """
        :param annotations_file_path: path to the directory containing directory
        'blobs_full' (which contains .npz files)
        :param labels_file_path: string with path to the file containing csv definition 
        of the dataset, default '{annotations_file_path}/cmb_blob_labels.csv', this 
        file has to contain columns 'ligands' and 'blob_map_file'
        """
        self.annotations_file_path = annotations_file_path
        if labels_file_path is None: 
            labels_file_path = os.path.join(self.annotations_file_path, 'cmb_blob_labels.csv')
        
        file_ligand_map = pd.read_csv(
            labels_file_path,
            usecols=['ligand', 'blob_map_filename']
        )
        self.file_ligand_map = file_ligand_map.set_index('blob_map_filename').to_dict()['ligand']
        self.files = list(self.file_ligand_map.keys())
        self.labels = list(self.file_ligand_map.values())

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        idx = self.files[idx]
        blob_path = os.path.join(self.annotations_file_path, 'blobs_full', idx)
        blob = np.load(blob_path)['blob']
        label = tensor(np.sum(blob), dtype=torch.float32)
        blob = tensor(blob, dtype=torch.float32)
        return (blob, label)

class DataLoader:
    """A class to represent simple dataloader which doesn't perform batching."""
    def __init__(self, dataset: Dataset):
        """
        :param dataset: dataset to be loaded
        """
        self.iter = iter(dataset)

    def __iter__(self):
        return self.iter

### to minkowski tensor ###

def to_minkowski_tensor(
    batch: torch.tensor
) -> ME.SparseTensor: 
    """
    Converts torch tensor containing blob or batch of blobs into MinkowskiEngine sparse tensor.
    :param batch: torch tensor
    :return: MinkowskiEngine sparse tensor
    """
    coordinates = torch.nonzero(batch).int()
    features = []
    for idx in coordinates:
        features.append(batch[tuple(idx)])
    features = torch.tensor(features).unsqueeze(-1)
    coordinates, features = ME.utils.sparse_collate([coordinates], [features])
    coordinates, features= ME.utils.sparse_quantize(coordinates, features)
    return ME.SparseTensor(features=features, coordinates=coordinates)

### NN ###

class NN(nn.Module):
    def __init__(
        self
    ):
        nn.Module.__init__(self)
        self.linear = nn.Linear(1, 1)

    def forward(
        self, 
        x: tensor
    ):
        x = x.F.sum()
        x = self.linear(x.unsqueeze(0))
        return x


class MinkowskiNN(ME.MinkowskiNetwork):
    def __init__(
        self
    ):
        ME.MinkowskiNetwork.__init__(self, 3)

        self.linear = ME.MinkowskiLinear(
            in_features=1, 
            out_features=1,
            bias=False
        )
        self.pool = ME.MinkowskiGlobalSumPooling()

    def forward(self, x: ME.SparseTensor):
        x = self.linear(x)
        x = self.pool(x)
        return x.F.squeeze(0).squeeze(0)


### trainer ###

model = MinkowskiNN()
dataset = LigandDataset('data')
dataloader = DataLoader(dataset)
criterion = torch.nn.L1Loss()
optimizer = torch.optim.Adam(
    model.parameters(),
    lr=1e-3
)

for idx, (blob, y) in enumerate(dataloader):
    if idx >= 10000:
        break
        
    optimizer.zero_grad()
    blob = to_minkowski_tensor(blob)
    y_hat = model(blob) # forward pass
    loss = criterion(y, y_hat)
    loss.backward()
    optimizer.step()

    if not idx % 100:
        #print([param for param in model.parameters()]) 
        print(f'iteration:{idx:>8}', f'loss: {loss.item():.4f}', f'groundtruth: {y.item():6.4f}', f'prediction: {y_hat.item():.4f}')
    
    if not idx % 500:
        print([param for param in model.parameters()]) 

print([param for param in model.parameters()])