Tutorial 8: Train NicheTrans on human lymph node 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_human_lymph_node import Lymph_node

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

warnings.filterwarnings("ignore")

Initialize the args and fix seeds

%run ./args/args_human_lymph_node.py
args = args

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

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

==========
Args:Namespace(adata_path='/home/wzk/ST_data/2024_nmethods_SpatialGlue_Human_lymph_node_3slides/', dropout_rate=0.1, eval_step=1, gamma=0.1, gpu_devices='0', lr=0.0003, max_epoch=40, n_source=3000, noise_rate=0.5, optimizer='adam', save_dir='./log', seed=1, stepsize=20, test_batch=32, train_batch=32, weight_decay=0.0005, workers=4)
==========

Initialize dataloaders and NicheTrans

# create the dataloaders
dataset = Lymph_node(adata_path=args.adata_path, n_top_genes=args.n_source)
trainloader, testloader = human_node_dataloader(args, dataset)

# create the model
source_dimension, target_dimension = dataset.rna_length, dataset.msi_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 13638 edges, 3484 cells.
3.9145 neighbors per cell on average.
------Calculating spatial graph...
The graph contains 27174 edges, 3484 cells.
7.7997 neighbors per cell on average.
------Calculating spatial graph...
The graph contains 13138 edges, 3359 cells.
3.9113 neighbors per cell on average.
------Calculating spatial graph...
The graph contains 26192 edges, 3359 cells.
7.7976 neighbors per cell on average.
=> SMA loaded
Dataset statistics:
  ------------------------------
  subset   | # num |
  ------------------------------
  train    |  After filting  3484 spots
  test     |  After filting  3359 spots
  ------------------------------

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_human_lymph_node_last.pth')
    ################

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