cnn_3d_vae_ddp.py

import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from train import make_subject_dict, get_dataloader
from datasets.fmri_4d import S1200
from tqdm import tqdm
from models.cnn_3d_vae import CNN_3d_AutoencoderKL
from ldm.modules.distributions.distributions import DiagonalGaussianDistribution
import torch.nn.functional as F
import torch.nn as nn
import torch.optim as optim
import numpy as np
import nibabel as nib
import matplotlib.pyplot as plt

# Jupyter 환경에서 inline 시각화

class temp():
    def __init__(self):
        self.distribution_type = 'ddp'
        self.rank = int(os.environ.get('RANK', 0))
        self.world_size = int(os.environ.get('WORLD_SIZE', 1))
        self.bs = 3
        self.workers = 2
        self.debug = False
        self.use_tb = True

def init_weights(m):
    if isinstance(m, nn.Conv3d) or isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight, nonlinearity='leaky_relu')
        if m.bias is not None:
            nn.init.constant_(m.bias, 0)
    elif isinstance(m, nn.Linear):
        nn.init.xavier_normal_(m.weight)
        if m.bias is not None:
            nn.init.constant_(m.bias, 0)

# DDP 초기화
def setup_ddp():
    # 마스터 정보 설정 (이미 torchrun으로 실행 시 환경변수 설정됨)
    dist.init_process_group(backend='nccl')
    local_rank = int(os.environ.get('LOCAL_RANK', 0))
    torch.cuda.set_device(local_rank)
    return local_rank

local_rank = setup_ddp()

train_metadata = f'/global/homes/s/seojw/4D-DiT/HCP_sex_metadata_full_train_set.csv'
test_metadata = f'/global/homes/s/seojw/4D-DiT/HCP_sex_metadata_rest_run2_RL_testset_debug.csv'

###########  Dataset & Dataloader ################
train_dict = make_subject_dict(train_metadata)
test_dict = make_subject_dict(test_metadata)
params = {
    'sequence_length':20, 
    'stride_within_seq':1,
    'stride_between_seq':1,
    'uint8_training':False,
    'padding':(0,0,0,0,0,0), 
    'root':'/home/'
}
opt = temp()
train_dataset = S1200(**params, subject_dict=train_dict, train=True)
test_dataset = S1200(**params, subject_dict=test_dict, train=False)

if opt.world_size > 1:
    # DistributedSampler 설정
    from torch.utils.data.distributed import DistributedSampler
    train_sampler = DistributedSampler(train_dataset, num_replicas=opt.world_size, rank=opt.rank, shuffle=True)
    test_sampler = DistributedSampler(test_dataset, num_replicas=opt.world_size, rank=opt.rank, shuffle=False)
else:
    train_sampler = None
    test_sampler = None

train_loader, test_loader, train_sampler_, test_sampler_ = get_dataloader(opt, train_dataset, test_dataset, train_sampler=train_sampler, test_sampler=test_sampler)

if opt.rank == 0:
    print(f'train_dataset = {len(train_dataset)}, test_dataset = {len(test_dataset)}')

#################################################
vae = CNN_3d_AutoencoderKL(
    double_z=True, z_channels=16, resolution=80, in_channels=1, out_ch=1, 
    ch=64, ch_mult=[1,2,4,4], num_res_blocks=2,
    attn_resolutions=[20, 10], dropout=0.0, embed_dim=16
).cuda(local_rank)
vae.apply(init_weights)

vae = DDP(vae, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True)

if opt.rank == 0:
    vae_num_params = sum(p.numel() for p in vae.parameters() if p.requires_grad)
    print(f"VAE - Number of trainable parameters: {vae_num_params / 1e6} M")

learning_rate = 1e-6
mae = nn.L1Loss()
optimizer = optim.Adam(vae.parameters(), lr=learning_rate)

for epoch in range(100):
    vae.train()
    train_loss = 0.0
    # DDP 사용 시 epoch 시작 시 sampler에 epoch 설정(Shuffle 일관성)
    if train_sampler is not None:
        train_sampler.set_epoch(epoch)

    for i, data in enumerate(train_loader):
        inputs = data['fmri_sequence'].type(torch.float).cuda(local_rank, non_blocking=True)
        volumes = inputs[:, :, :, :, :, 10]
        recon, posterior = vae(volumes)

        recon_loss = mae(recon, volumes)
        mean = posterior.mean
        logvar = posterior.logvar
        kl_loss = -0.5 * torch.mean(1 + logvar - mean.pow(2) - logvar.exp())
        loss = recon_loss + 1e-6 * kl_loss

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_loss += loss.item()

        if opt.rank == 0 and i % 20 == 0:
            print(f"iter {i}, loss: {loss.item()}")
            vae.eval()
            with torch.no_grad():
                sample_data = train_dataset[0]
                sample_input = sample_data['fmri_sequence'].unsqueeze(0).type(torch.float).to(local_rank)
                sample_volumes = sample_input[:, :, :, :, :, 10]
                recon_sample, _ = vae(sample_volumes)

            original_img = sample_volumes[0,0].cpu().numpy()
            recon_img = recon_sample[0,0].cpu().numpy()
            original_nii = nib.Nifti1Image(original_img, affine=np.eye(4))
            recon_nii = nib.Nifti1Image(recon_img, affine=np.eye(4))

            mid_z = 40
            fig, axes = plt.subplots(1, 2, figsize=(8,4))
            axes[0].imshow(original_img[:,:,mid_z].T, cmap='gray', origin='lower')
            axes[0].set_title('Original')
            axes[0].axis('off')

            axes[1].imshow(recon_img[:,:,mid_z].T, cmap='gray', origin='lower')
            axes[1].set_title('Reconstructed')
            axes[1].axis('off')

            plt.suptitle(f"Epoch {epoch}, iter {i}")
            plt.tight_layout()
            # GPU 번호(local_rank)를 파일명에 넣는다면:
            gpu_num = local_rank
            # 원하는 대로 저장
            if i % 400 == 0:
                nib.save(original_nii, f"/pscratch/sd/s/seojw/4D-DiT/recon/3DCNN/epoch_{epoch}_iter_{i}_original_gpu{gpu_num}.nii.gz")
                nib.save(recon_nii, f"/pscratch/sd/s/seojw/4D-DiT/recon/3DCNN/epoch_{epoch}_iter_{i}_recon_gpu{gpu_num}.nii.gz")
                fig.savefig(f"/pscratch/sd/s/seojw/4D-DiT/recon/3DCNN/epoch_{epoch}_iter_{i}_gpu{gpu_num}.png")
            plt.show()
            plt.close(fig)

    avg_train_loss = train_loss / len(train_loader)
    if opt.rank == 0:
        print(f"Epoch [{epoch}/{100}] - Avg Train Loss: {avg_train_loss:.4f}")

dist.barrier()
if opt.rank == 0:
    print("Training completed.")