Tutorial 5: Train NicheTrans on STARmap PLUS data

import os, time, datetime, warnings

import torch
import torch.nn as nn
from torch.optim import lr_scheduler

from model.nicheTrans_img import *
from datasets.data_manager_STARmap_PLUS import AD_Mouse

from utils.utils import *
from utils.utils_training_STARmap_PLUS import train, test
from utils.utils_dataloader import *

warnings.filterwarnings("ignore")

Initialize the args and fix seeds

%run ./args/args_STARmap_PLUS.py
args = args

set_seed(args.seed)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices

print("==========\nArgs:{}\n==========".format(args))

==========
Args:Namespace(AD_adata_path='/home/wzk/ST_data/AD_mouse2/norm/AD_mouses_adata', Wild_type_adata_path='/home/wzk/ST_data/AD_mouse2/norm/benigh_mouses', dropout_rate=0.25, eval_step=1, gamma=0.1, gpu_devices='0', label_path='/home/wzk/ST_data/AD_mouse/generated_by_zhikang_2_filtering_tao15_abeta50', lr=0.0001, max_epoch=20, n_top_genes=2000, noise_rate=0.5, optimizer='adam', save_dir='./log', seed=1, stepsize=20, test_batch=32, train_batch=128, weight_decay=0.0005, workers=4)
==========

Initialize dataloaders and NicheTrans

# create the dataloaders
dataset = AD_Mouse(AD_adata_path=args.AD_adata_path, Wild_type_adata_path=args.Wild_type_adata_path, label_path=args.label_path, n_top_genes=args.n_top_genes)
trainloader, testloader = ad_mouse_dataloader(args, dataset)

# create the model
source_dimension, target_dimension = dataset.rna_length, dataset.target_length
model = NicheTrans(source_length=source_dimension, target_length=target_dimension, noise_rate=args.noise_rate, dropout_rate=args.dropout_rate)
model = nn.DataParallel(model).cuda()

------Calculating spatial graph...
The graph contains 124464 edges, 10372 cells.
12.0000 neighbors per cell on average.
------Calculating spatial graph...
The graph contains 115608 edges, 9634 cells.
12.0000 neighbors per cell on average.
------Calculating spatial graph...
The graph contains 96408 edges, 8034 cells.
12.0000 neighbors per cell on average.
=> AD Mouse loaded
Dataset statistics:
  ------------------------------
  subset   | # num |
  ------------------------------
  train    |  10372 spots, 894.0 positive tao, 291.0 positive plaque
  test     |   9634 spots, 620.0 positive tao, 195.0 positive plaque
  ------------------------------

Initialize loss function (criterion) and optimizer

criterion = nn.MSELoss()

if args.optimizer == 'adam':
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
elif args.optimizer == 'SGD':
    optimizer = torch.optim.SGD(model.parameters(), lr=args.lr)
else:
    print('unexpected optimizer')

if args.stepsize > 0:
    scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma)

Model training and testing

start_time = time.time()

for epoch in range(args.max_epoch):
    last_epoch = epoch + 1 == args.max_epoch

    print("==> Epoch {}/{}".format(epoch+1, args.max_epoch))

    ################
    train(args, model, criterion, optimizer, trainloader, dataset.target_panel)
    if args.stepsize > 0: scheduler.step()

    if (epoch+1) % args.eval_step == 0:
        pearson = test(args, model, testloader, dataset.target_panel, last_epoch)

    if last_epoch==True:
        torch.save(model.state_dict(), 'NicheTrans_STARmap_PLUS.pth')
    ################

elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))