use burn::lr_scheduler::cosine::CosineAnnealingLrSchedulerConfig;
use burn::module::Module;
use burn::nn::loss::CrossEntropyLossConfig;
use burn::nn::pool::{MaxPool1d, MaxPool1dConfig};
use burn::nn::{Linear, LinearConfig, PaddingConfig1d, Relu};
use burn::nn::{conv::Conv1d, conv::Conv1dConfig};
use burn::optim::AdamConfig;
use burn::record::CompactRecorder;
use burn::tensor::{Int, Tensor, backend::Backend};
use burn::train::metric::{AccuracyMetric, LossMetric};
use burn::train::{
ClassificationOutput, InferenceStep, Learner, SupervisedTraining, TrainOutput, TrainStep,
};
use burn::{
config::Config,
data::{dataloader::batcher::Batcher, dataset::Dataset},
};
use std::error::Error;
use std::fs::File;
use std::io::{BufRead, BufReader};
const SEQ_LEN: usize = 15;
const DATA_PATH: &str = "data.csv";
pub fn readfile(path: &str) -> Result<Vec<(String, usize)>, Box<dyn Error>> {
let file = File::open(path).expect("file not found");
let fileread = BufReader::new(file);
let mut veclabel: Vec<(String, usize)> = Vec::new();
for i in fileread.lines() {
let line = i.expect("readfile not present");
let linevec = line.split(",").collect::<Vec<_>>();
veclabel.push((linevec[0].to_string(), linevec[1].parse::<usize>().unwrap()));
}
Ok(veclabel)
}
#[derive(Debug, Clone)]
pub struct DnaItem {
pub seq: String,
pub label: usize,
}
#[derive(Debug, Clone)]
pub struct DnaDataSet {
items: Vec<DnaItem>,
}
impl DnaDataSet {
fn train(pathname: &str) -> Self {
let raw = readfile(pathname).unwrap();
Self {
items: raw
.into_iter()
.map(|(s, l)| DnaItem { seq: s, label: l })
.collect(),
}
}
}
impl Dataset<DnaItem> for DnaDataSet {
fn get(&self, index: usize) -> Option<DnaItem> {
self.items.get(index).cloned()
}
fn len(&self) -> usize {
self.items.len()
}
}
pub fn onehot_encode(seq: &str) -> Vec<f32> {
let seq = seq.to_uppercase();
let mut encoded = vec![0f32; 4 * SEQ_LEN];
for (pos, base) in seq.chars().enumerate().take(SEQ_LEN) {
let channel = match base {
'A' => 0,
'C' => 1,
'G' => 2,
'T' => 3,
_ => continue,
};
encoded[channel * SEQ_LEN + pos] = 1.0;
}
encoded
}
#[derive(Clone, Debug, Default)]
pub struct DnaBatcher<B: Backend> {
_backend: std::marker::PhantomData<B>,
}
#[derive(Debug, Clone)]
pub struct DnaBatch<B: Backend> {
input: Tensor<B, 3>,
targets: Tensor<B, 1, Int>,
}
impl<B: Backend> Batcher<B, DnaItem, DnaBatch<B>> for DnaBatcher<B> {
fn batch(&self, items: Vec<DnaItem>, device: &B::Device) -> DnaBatch<B> {
let inputs = items
.iter()
.map(|it| {
let flat = onehot_encode(&it.seq);
Tensor::<B, 1>::from_floats(flat.as_slice(), device).reshape([4, SEQ_LEN])
})
.collect::<Vec<_>>();
let inputs = Tensor::stack(inputs, 0);
let targets = items
.iter()
.map(|it| Tensor::<B, 1, Int>::from_ints([it.label as i32], device))
.collect::<Vec<_>>();
let targets = Tensor::cat(targets, 0);
DnaBatch {
input: inputs,
targets,
}
}
}
#[derive(Module, Debug)]
pub struct DnaCnnClassifier<B: Backend> {
conv1: Conv1d<B>,
pool1: MaxPool1d,
conv2: Conv1d<B>,
pool2: MaxPool1d,
fc: Linear<B>,
activation: Relu,
flattened_size: usize,
}
#[derive(Config, Debug)]
pub struct DnaCnnClassifierConfig {
#[config(default = 16)]
conv1_channels: usize,
#[config(default = 32)]
conv2_channels: usize,
#[config(default = 3)]
kernel_size: usize,
#[config(default = 2)]
pool_size: usize,
}
impl DnaCnnClassifierConfig {
fn init<B: Backend>(&self, device: &B::Device) -> DnaCnnClassifier<B> {
let conv1 = Conv1dConfig::new(4, self.conv1_channels, self.kernel_size)
.with_padding(PaddingConfig1d::Same)
.init(device);
let conv2 = Conv1dConfig::new(self.conv1_channels, self.conv2_channels, self.kernel_size)
.with_padding(PaddingConfig1d::Same)
.init(device);
let pool1 = MaxPool1dConfig::new(self.pool_size).init();
let pool2 = MaxPool1dConfig::new(self.pool_size).init();
let len_after_pool1 = SEQ_LEN / self.pool_size;
let len_after_pool2 = len_after_pool1 / self.pool_size;
let flattened_size = self.conv2_channels * len_after_pool2;
DnaCnnClassifier {
conv1,
pool1,
conv2,
pool2,
fc: LinearConfig::new(flattened_size, 2).init(device),
activation: Relu::new(),
flattened_size,
}
}
}
impl<B: Backend> DnaCnnClassifier<B> {
fn forward(&self, input: Tensor<B, 3>) -> Tensor<B, 2> {
let x = self.conv1.forward(input);
let x = self.activation.forward(x);
let x = self.pool1.forward(x);
let x = self.conv2.forward(x);
let x = self.activation.forward(x);
let x = self.pool2.forward(x);
let batch_size = x.dims()[0];
let x = x.reshape([batch_size, self.flattened_size]);
self.fc.forward(x)
}
fn forward_classification(&self, batch: DnaBatch<B>) -> ClassificationOutput<B> {
let logits = self.forward(batch.input);
let loss = CrossEntropyLossConfig::new()
.init(&logits.device())
.forward(logits.clone(), batch.targets.clone());
ClassificationOutput {
loss,
output: logits,
targets: batch.targets,
}
}
}
impl<B: burn::tensor::backend::AutodiffBackend> TrainStep for DnaCnnClassifier<B> {
type Input = DnaBatch<B>;
type Output = ClassificationOutput<B>;
fn step(&self, batch: DnaBatch<B>) -> TrainOutput<ClassificationOutput<B>> {
let item = self.forward_classification(batch);
TrainOutput::new(self, item.loss.backward(), item)
}
}
impl<B: Backend> InferenceStep for DnaCnnClassifier<B> {
type Input = DnaBatch<B>;
type Output = ClassificationOutput<B>;
fn step(&self, batch: DnaBatch<B>) -> ClassificationOutput<B> {
self.forward_classification(batch)
}
}
pub fn train<B: burn::tensor::backend::AutodiffBackend>(device: B::Device, datapath: &str) {
let dataset = DnaDataSet::train(datapath);
let batcher = DnaBatcher::<B>::default();
let dataloader_train = burn::data::dataloader::DataLoaderBuilder::new(batcher)
.batch_size(2)
.shuffle(42)
.build(dataset);
let dataloader_valid =
burn::data::dataloader::DataLoaderBuilder::new(DnaBatcher::<B::InnerBackend>::default())
.batch_size(2)
.build(DnaDataSet::train(DATA_PATH));
let model = DnaCnnClassifierConfig::new().init(&device);
let optimizer = AdamConfig::new().init();
let num_epochs = 20;
let batch_size = 2;
let dataset_len = DnaDataSet::train(DATA_PATH).len();
let batches_per_epoch = (dataset_len + batch_size - 1) / batch_size;
let num_iters = num_epochs * batches_per_epoch;
let lr_scheduler = CosineAnnealingLrSchedulerConfig::new(1e-3, num_iters)
.init()
.expect("lr scheduler config should be valid");
let training = SupervisedTraining::new("./dna_huget/model", dataloader_train, dataloader_valid)
.metrics((AccuracyMetric::new(), LossMetric::new()))
.with_file_checkpointer(CompactRecorder::new())
.num_epochs(num_epochs)
.summary();
let result = training.launch(Learner::new(model, optimizer, lr_scheduler));
result
.model
.save_file("./dna_huget/model", &CompactRecorder::new())
.expect("model should save");
}