train.py

#coding=utf-8
import torch
from torch.autograd import Variable
from torch.backends import cudnn
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import numpy as np
import pprint
import transforms
from dataloader_fft import KFDataset
#from models import KFSGNet
import os
import argparse
import matplotlib.pyplot as plt
from loss import KpLoss,CLALoss
import tempfile
# existing code...
import torch.distributed as dist
import glob as _glob
import json
from U2Netnew import U2Net  # or your actual U2Net import

def reduce_value(value, average=True):
    """Safe reduce: if no distributed init, return value unchanged."""
    if not (dist.is_available() and dist.is_initialized()):
        return value
    if isinstance(value, torch.Tensor):
        dist.reduce(value, dst=0)
        if average and dist.get_world_size() > 0:
            value = value / dist.get_world_size()
        return value
    return value
# existing code...
config = dict()
config['lr'] = 0.01
config['momentum'] = 0.009
config['weight_decay'] = 1e-4
config['epoch_num'] = 100
config['batch_size'] = 2
config['sigma'] = 2.5
config['debug_vis'] = False

config['train_fname'] = ''
config['test_fname'] =''
#config ['path_image'] = '/public/huangjunzhang/KeyPointsDetection-master/dataloader_train/'
config ['test_image_path'] = '/kaggle/input/vindr-spinexray/physionet.org/files/vindr-spinexr/1.0.0/test_images'
config ['train_image_path'] = '/kaggle/input/annotated-medical-image-dataset-for-spinal-lesions/physionet.org/files/vindr-spinexr/1.0.0/train_images'

# config ['test_image_path'] = '/public/huangjunzhang/KeyPointsDetection-master/lumbar_test/'
# config ['train_image_path'] = '/public/huangjunzhang/KeyPointsDetection-master/lumbar_train/'

config['path_label'] = '/public/huangjunzhang/KeyPointsDetection-master/txt/'
#config['path_label_train'] = '/public/huangjunzhang/KeyPointsDetection-master/txt/train_json/'
config['path_label_train'] = '/public/huangjunzhang/KeyPointsDetection-master/txt/lumbar_json'
#config['json_path']='/public/huangjunzhang/test/keypoints_train.json'
config['is_test'] = False

config['save_freq'] = 10
config['checkout'] = '/public/huangjunzhang/KeyPointsDetection-master/Checkpoints/kd_MLT_epoch_499_model.ckpt'
config['start_epoch'] = 0
config['load_pretrained_weights'] = False
config['eval_freq'] = 50
config['debug'] = False
config['featurename2id'] = {
    'C2_TR':0,
    'C2_TL':1,
    'C2_DR':2,
    'C2_DL':3,
    'C3_TR':4,
    'C3_TL':5,
    'C3_DR':6,
    'C3_DL':7,
    'C4_TR': 8,
    'C4_TL': 9,
    'C4_DR': 10,
    'C4_DL': 11,
    'C5_TR': 12,
    'C5_TL': 13,
    'C5_DR': 14,
    'C5_DL': 15,
    'C6_TR': 16,
    'C6_TL': 17,
    'C6_DR': 18,
    'C6_DL': 19,
    'C7_TR': 20,
    'C7_TL': 21,
    'C7_DR': 22,
    'C7_DL': 23,
    }


def get_peak_points(heatmaps):
    """

    :param heatmaps: numpy array (N,24,256,256)
    :return:numpy array (N,24,2) #
    """
    N,C,H,W = heatmaps.shape   # N= batch size C=24 hotmaps
    all_peak_points = []
    for i in range(N):
        peak_points = []
        for j in range(C):
            yy,xx = np.where(heatmaps[i,j] == heatmaps[i,j].max())
            nan = float('nan')
            #print(yy)
            y = yy[0] if yy[0]!= None else 0
            x = xx[0]
            peak_points.append([x,y])
        all_peak_points.append(peak_points)
    all_peak_points = np.array(all_peak_points)
    return all_peak_points

def get_mse(pred_points,gts,indices_valid=None):
    """

    :param pred_points: numpy (N,4,2)
    :param gts: numpy (N,4,2)
    :return:
    """
    pred_points = pred_points[indices_valid[0],indices_valid[1],:]
    gts = gts[indices_valid[0],indices_valid[1],:]
    pred_points = Variable(torch.from_numpy(pred_points).float(),requires_grad=False)
    gts = Variable(torch.from_numpy(gts).float(),requires_grad=False)
    criterion = nn.MSELoss()
    loss = criterion(pred_points,gts)
    return loss

def calculate_mask(heatmaps_targets):
    """

    :param heatmaps_target: Variable (N,4,256,256)
    :return: Variable (N,4,256,256)
    """
    N,C,_,_ = heatmaps_targets.size()  #N =8 C = 4
    N_idx = []
    C_idx = []
    for n in range(N):      # 0-7
        for c in range(C):  # 0-3
            max_v = heatmaps_targets[n,c,:,:].max().item()
            if max_v != 0.0:
                N_idx.append(n)
                C_idx.append(c)
    mask = Variable(torch.zeros(heatmaps_targets.size()))
    mask[N_idx,C_idx,:,:] = 1.0
    mask = mask.float().cuda()
    return mask,[N_idx,C_idx]

def main(args):
    if torch.cuda.is_available() is False:
        raise EnvironmentError("not find GPU device for training.")

     # Auto-detect label folder if config path missing (Kaggle / mounted inputs)
    # if not os.path.exists(config.get('path_label_train', '')):
    #     cand = _glob.glob('/kaggle/input/**/**/*_jpg_Label.json', recursive=True)
    #     if not cand:
    #         cand = _glob.glob('/kaggle/input/**/*_jpg_Label.json', recursive=True)
    #     if cand:
    #         config['path_label_train'] = os.path.dirname(cand[0])
    #         config['path_label'] = config['path_label_train']
    #         print(f"Detected label directory: {config['path_label_train']}")
    #     else:
    #         print("ERROR: no label json files found under /kaggle/input.")
    #         print("Set config['path_label_train'] to the directory that contains *_jpg_Label.json files.")
    #         return
         # If label folder is missing, try to find jsons next to images; otherwise generate placeholder jsons
    if not os.path.exists(config.get('path_label_train', '')):
        # try to find existing json labels under the train image path
        candidates = _glob.glob(os.path.join(config['train_image_path'], '**', '*.json'), recursive=True)
        if candidates:
            config['path_label_train'] = os.path.dirname(candidates[0])
            config['path_label'] = config['path_label_train']
            print(f"Detected label directory: {config['path_label_train']}")
        else:
            # collect image files
            imgs = []
            for ext in ('*.jpg', '*.jpeg', '*.png', '*.dcm', '*.DCM', '*.dicom'):
                imgs += _glob.glob(os.path.join(config['train_image_path'], '**', ext), recursive=True)
            if not imgs:
                print("ERROR: no images found at", config['train_image_path'])
                print("Set config['train_image_path'] and config['path_label_train'] correctly and re-run.")
                return
            # create writable placeholder label folder and JSONs
            gen_dir = os.path.join(os.getcwd(), 'generated_labels')
            os.makedirs(gen_dir, exist_ok=True)
            for p in imgs:
                name = os.path.splitext(os.path.basename(p))[0] + '.json'
                jpath = os.path.join(gen_dir, name)
                if not os.path.exists(jpath):
                    with open(jpath, 'w') as f:
                        # minimal placeholder structure expected by many loaders
                        json.dump({'image': os.path.basename(p), 'annotations': []}, f)
            config['path_label_train'] = gen_dir
            config['path_label'] = gen_dir
            print("Generated placeholder labels in", gen_dir)
    # ...existing code...
      
    # existing code...
    # rank = args.rank
    # device = torch.device(args.device)
    # batch_size = args.batch_size
    #weights_path = args.weights
    # args.lr *= args.world_size  # 学习率要根据并行GPU的数量进行倍增
    # provide safe defaults when running single-process (no --rank/--gpu provided)
    rank = getattr(args, 'rank', 0)
    args.rank = rank
    args.gpu = getattr(args, 'gpu', 0)
    device = torch.device(args.device)
    batch_size = args.batch_size
    # adjust lr only if world_size provided
    if hasattr(args, 'world_size') and args.world_size is not None:
        args.lr *= args.world_size  # 学习率要根据并行GPU的数量进行倍增
# ...existing code..
   
    pprint.pprint(config)

    data_transforms = {
        "train": transforms.Compose([
                                     transforms.RandomHorizontalFlip(0.5),
                                     transforms.Blur(),
                                     transforms.Brightness(),
                                     transforms.ToTensor(),
                                     # transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
                                     ]),

        "val" : transforms.Compose([transforms.ToTensor(),
                                    # transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
                                    ])
    }


    trainDataset = KFDataset(config , mode='train', transforms=data_transforms["train"])
    # train_sampler = torch.utils.data.distributed.DistributedSampler(trainDataset)
    testDataset = KFDataset(config, mode='test',transforms=data_transforms["val"])
    # test_sampler = torch.utils.data.distributed.DistributedSampler(testDataset)
    # Skip training if dataset is empty (paths may not exist in this environment)
    if len(trainDataset) == 0:
        print(f"⚠️  trainDataset is empty (0 samples). Check config paths:")
        print(f"  train_image_path: {config['train_image_path']}")
        print(f"  path_label_train: {config['path_label_train']}")
        print("Skipping training.")
        return

    if dist.is_available() and dist.is_initialized():
        train_sampler = torch.utils.data.distributed.DistributedSampler(trainDataset)
        test_sampler = torch.utils.data.distributed.DistributedSampler(testDataset)
    else:
        from torch.utils.data import RandomSampler, SequentialSampler
        train_sampler = RandomSampler(trainDataset)
        test_sampler = SequentialSampler(testDataset)
# ...existing code...
    #testDataLoader = DataLoader(testDataset,1, True, num_workers=8)
    # 将样本索引每batch_size个元素组成一个list
    train_batch_sampler = torch.utils.data.BatchSampler(
        train_sampler, batch_size, drop_last=True)

    nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers
    if rank == 0:
        print('Using {} dataloader workers every process'.format(nw))
    #trainDataset.load()
    # 定义 data loader
    sample_num = len(trainDataset)
    print(sample_num)



    trainDatasetloader = torch.utils.data.DataLoader(trainDataset,
                                               batch_sampler=train_batch_sampler,
                                               pin_memory=True,
                                               num_workers=nw)
                                               #collate_fn=trainDataset.collate_fn)
    testDataLoader = torch.utils.data.DataLoader(testDataset,
                                               batch_sampler=test_sampler,
                                               pin_memory=True,
                                               num_workers=nw)
                                               #collate_fn=trainDataset.collate_fn)



    # collate_fn=trainDataset.collate_fn)

    torch.manual_seed(0)
    cudnn.benchmark = True
    #model = KFSGNet()
    model = U2Net(in_channels=1,out_channels=24)
    model = model.float().cuda().to(device)
    if config['load_pretrained_weights']:
         if (config['checkout'] != ''):
             print("load dict from checkpoint")
             model.load_state_dict(torch.load(config['checkout']))

    else:
        checkpoint_path = os.path.join(tempfile.gettempdir(),"initial_weights.pt")
        if rank ==0 :
            torch.save(model.state_dict(),checkpoint_path)
        if dist.is_available() and dist.is_initialized():
            dist.barrier()

        model.load_state_dict(torch.load(checkpoint_path,map_location=device))
    # convert to DDP model
    # if args.syncBN:
    #     model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)

    # model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
    # convert to DDP model only when distributed is initialized
    if dist.is_available() and dist.is_initialized():
        if args.syncBN:
            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
    else:
        # single-process: keep model on device
        model = model.to(device)

    pg = [p for p in model.parameters() if p.requires_grad]
    optimizer = optim.SGD(pg, lr=args.lr, momentum=0.9, weight_decay=0.005)
    #print('test')

    #criterion = nn.MSELoss(reduction='sum')
    #criterion = nn.BCELoss()
    criterion = KpLoss()
    class_criterion = CLALoss()
    train_loss = []
    val_loss = []
    for epoch in range(config['start_epoch'],config['epoch_num']+config['start_epoch']):
        train_sampler.set_epoch(epoch)
        model.train()
        for i, (inputs, heatmaps_targets, gts, loss_mask, label ) in enumerate(trainDatasetloader):

            lam = 0.01
            inputs = Variable(inputs).cuda().float()

            heatmaps_targets = Variable(heatmaps_targets).cuda()
            mask, indices_valid = calculate_mask(heatmaps_targets)

            optimizer.zero_grad()
            outputs, class_output = model(inputs)
            outputs = outputs.to(torch.float32)

            heatmaps_targets = heatmaps_targets.to(torch.float32)
            # print(torch.max(outputs[0]), torch.min(outputs[0]))
            # print(torch.max(heatmaps_targets[0]),torch.min(heatmaps_targets[0]))
            outputs = outputs * mask
            heatmaps_targets = heatmaps_targets * mask

            kp_loss = criterion(outputs, heatmaps_targets, loss_mask)
            class_loss = class_criterion(class_output, label, loss_mask)
            running_loss = kp_loss + class_loss
            running_loss.backward()
            running_loss = reduce_value(running_loss, average=True)
            optimizer.step()

            # 统计最大值与最小值
            v_max = torch.max(outputs)
            v_min = torch.min(outputs)
            if (i + 1) % config['eval_freq'] == 0:
                print('---------------calculate loss-------')
                print('[ Epoch {:005d} -> {:005d} / {} ] loss : {:15} CLASSLOSS:{} max : {:10} min : {}'.format(
                    epoch, i * config['batch_size'],
                    sample_num, running_loss.item(), class_loss.item(), v_max.item(), v_min.item()))

        train_loss.append(running_loss)

        if (epoch+1) % config['save_freq'] == 0 or epoch == config['epoch_num'] - 1:
            torch.save(model.state_dict(),'./Checkpoints/kd_MLTGPU_epoch_{}_model.ckpt'.format(epoch))
            

    plt.figure()
    plt.plot(train_loss, 'b-', label='Recon_loss')
    plt.ylabel('Train_loss')
    plt.xlabel('iter_num')
    plt.show()

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--epochs', type=int, default=config['epoch_num'])
    parser.add_argument('--batch-size', type=int, default=config['batch_size'])
    parser.add_argument('--lr', type=float, default=config['lr'])
    parser.add_argument('--lrf', type=float, default=0.1)
    # 是否启用SyncBatchNorm
    parser.add_argument('--syncBN', type=bool, default=True)
    parser.add_argument('--device', default='cuda', help='device id (i.e. 0 or 0,1 or cpu)')
    # 开启的进程数(注意不是线程),不用设置该参数，会根据nproc_per_node自动设置
    parser.add_argument('--world-size', default=4, type=int,
                        help='number of distributed processes')
    parser.add_argument('--dist-url', default='env://', help='url used to set up distributed training')
    opt = parser.parse_args()
    main(opt)