Upload modle

config.json

@@ -0,0 +1,27 @@
+{
+  "architectures": [
+    "ViTMAEForPreTraining"
+  ],
+  "attention_probs_dropout_prob": 0.0,
+  "decoder_hidden_size": 512,
+  "decoder_intermediate_size": 1024,
+  "decoder_num_attention_heads": 16,
+  "decoder_num_hidden_layers": 8,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.0,
+  "hidden_size": 768,
+  "image_size": 112,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-12,
+  "mask_ratio": 0.75,
+  "model_type": "vit_mae",
+  "norm_pix_loss": false,
+  "num_attention_heads": 12,
+  "num_channels": 1,
+  "num_hidden_layers": 12,
+  "patch_size": 8,
+  "qkv_bias": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.41.2"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:13d8f408d176af86d658b98d009de423978dc6969adb5e63f5447ee8c99982db
+size 410476120

pretraining_pl_DDP_v5.py

@@ -0,0 +1,213 @@
+import csv
+import h5py
+import torch
+import torch.nn as nn
+import random
+import numpy as np
+import os
+import shutil
+import pandas as pd
+from torchvision import transforms
+from PIL import Image
+from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler, Subset, random_split
+import torch.optim as optim
+import time
+from tqdm import tqdm
+from torch.optim import lr_scheduler
+from transformers import ViTFeatureExtractor, AutoImageProcessor, ViTMAEConfig, ViTMAEModel, ViTMAEForPreTraining
+from torchvision.datasets import ImageFolder
+import lightning.pytorch as pl
+from lightning.pytorch import Trainer
+from lightning.pytorch.callbacks import ModelCheckpoint, LearningRateMonitor, RichProgressBar
+from lightning.pytorch.loggers import TensorBoardLogger
+from lightning.pytorch.callbacks import RichProgressBar
+from lightning.pytorch.callbacks import TQDMProgressBar
+from lightning.pytorch.utilities import rank_zero_only
+
+DEVICE_NUM = torch.cuda.device_count()
+os.environ['CUDA_VISIBLE_DEVICES'] = ','.join([str(i) for i in range(DEVICE_NUM)])
+
+SEED = 42
+DATA_DIR = "../../0.data/pretrain_nucleus_image_all_16M.hdf5"
+BATCH_SIZE = 400 *2
+NUM_EPOCHS = 70
+LEARNINGRATE = 0.0001
+PROJECT_NAME = 'Nuspire_Pretraining_V5'
+
+transform = transforms.Compose([
+    transforms.Grayscale(),
+    transforms.RandomResizedCrop((112, 112), scale=(0.5625, 1.0), ratio=(0.75, 1.33)),
+    transforms.RandomHorizontalFlip(p=0.5),
+    transforms.RandomVerticalFlip(p=0.5),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.21869252622127533], std=[0.1809280514717102])
+])
+
+configuration = ViTMAEConfig(
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        intermediate_size=3072,
+        hidden_act="gelu",
+        hidden_dropout_prob=0.0,
+        attention_probs_dropout_prob=0.0,
+        initializer_range=0.02,
+        layer_norm_eps=1e-12,
+        image_size=112,
+        patch_size=8,
+        num_channels=1,
+        qkv_bias=True,
+        decoder_num_attention_heads=16,
+        decoder_hidden_size=512,
+        decoder_num_hidden_layers=8,
+        decoder_intermediate_size=1024,
+        mask_ratio=0.75,
+        norm_pix_loss=False
+)
+
+class HDF5Dataset(Dataset):
+    def __init__(self, hdf5_path, transform=None):
+        self.hdf5_path = hdf5_path
+        self.transform = transform
+        self.hdf5_file = h5py.File(hdf5_path, 'r', rdcc_nbytes=10*1024**3, rdcc_w0=0.0, rdcc_nslots=10007)
+        self.images = self.hdf5_file['images']
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        img = self.images[idx]
+        
+        if self.transform:
+            img = Image.fromarray(img)
+            img = self.transform(img)
+        
+        return img
+
+    def __del__(self):
+        self.hdf5_file.close()
+
+class NucleusDataModule(pl.LightningDataModule):
+    def __init__(self, dataset, batch_size):
+        super().__init__()
+        self.dataset = dataset
+        self.batch_size = batch_size
+
+    def setup(self, stage=None):
+        
+        train_size = int(0.8 * len(self.dataset))
+        test_size = len(self.dataset) - train_size
+        self.train_dataset, self.test_dataset = random_split(self.dataset, [train_size, test_size])
+
+    def train_dataloader(self):
+        return DataLoader(self.train_dataset, batch_size=self.batch_size, num_workers=16, pin_memory=True, prefetch_factor=5)
+
+    def val_dataloader(self):
+        return DataLoader(self.test_dataset, batch_size=self.batch_size * 3, num_workers=16, pin_memory=True, prefetch_factor=5)
+
+class ViTMAEPreTraining(pl.LightningModule):
+    def __init__(self, configuration):
+        super().__init__()
+        self.model = ViTMAEForPreTraining(configuration) 
+        self.save_hyperparameters()
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        x = batch
+        x = x.to(self.device)
+        outputs = self.model(x)
+        loss = outputs.loss
+        self.log('train_loss', loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True)
+        return loss
+
+   
+    def validation_step(self, batch, batch_idx):
+        x = batch
+        x = x.to(self.device)
+        outputs = self.model(x)
+        loss = outputs.loss
+        self.log('val_loss', loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True)
+        return loss       
+    
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.model.parameters(), lr=LEARNINGRATE)
+        warmup_epochs = 10
+        warmup_factor = lambda epoch: epoch / warmup_epochs if epoch < warmup_epochs else 1
+        scheduler_warmup = torch.optim.lr_scheduler.LambdaLR(optimizer, warmup_factor)
+        scheduler_regular = torch.optim.lr_scheduler.StepLR(optimizer, 20, gamma=0.5)
+        scheduler = {
+            'scheduler': torch.optim.lr_scheduler.SequentialLR(optimizer, schedulers=[scheduler_warmup, scheduler_regular], milestones=[warmup_epochs]),
+            'interval': 'epoch',
+            'frequency': 1
+        }
+        return [optimizer], [scheduler]
+
+class EpochLoggingCallback(pl.Callback):
+    def __init__(self):
+        super().__init__()
+    
+    @rank_zero_only
+    def on_validation_epoch_end(self, trainer, pl_module):
+        train_loss = trainer.callback_metrics.get('train_loss')
+        val_loss = trainer.callback_metrics.get('val_loss')
+        if train_loss is not None and val_loss is not None:
+            trainer.logger.experiment.add_scalars(
+                "Epoch/Loss",
+                {'Train Loss': train_loss, 'Validation Loss': val_loss},
+                trainer.current_epoch
+            )
+
+class SaveEpochModelCallback(pl.Callback):
+    def __init__(self):
+        super().__init__()
+    
+    @rank_zero_only
+    def on_validation_epoch_end(self, trainer, pl_module):
+        path = trainer.checkpoint_callback.dirpath
+        epoch = trainer.current_epoch
+        pl_module.model.save_pretrained(f'{path}/epoch{epoch}')
+
+dataset = HDF5Dataset(hdf5_path=DATA_DIR, transform=transform)
+
+data_module = NucleusDataModule(dataset, BATCH_SIZE)
+
+epoch_logging_callback = EpochLoggingCallback()
+
+save_epoch_model_callback = SaveEpochModelCallback()
+
+progress_bar = RichProgressBar()
+
+logger = TensorBoardLogger(save_dir=f'./{PROJECT_NAME}_outputs', name="tensorboard")
+
+best_model_callback = ModelCheckpoint(
+    dirpath=f'./{PROJECT_NAME}_outputs/model',
+    filename='{epoch:02d}-{val_loss:.2f}',
+    save_top_k=3,
+    mode='min',
+    monitor='val_loss'
+)
+
+lr_monitor = LearningRateMonitor(logging_interval='epoch')
+
+trainer = Trainer(
+    max_epochs=NUM_EPOCHS,
+    devices=DEVICE_NUM,  # 设置使用的设备数量
+    accelerator='gpu',  # 指定使用GPU
+    strategy='ddp',
+    logger=logger,
+    callbacks=[lr_monitor,
+               progress_bar, 
+               epoch_logging_callback, 
+               save_epoch_model_callback, 
+               best_model_callback]
+)
+
+# 设置随机种子
+pl.seed_everything(SEED, workers=True)
+
+model = ViTMAEPreTraining(configuration,)
+trainer.fit(model, data_module)
+
+