diff --git a/pretrain_ict.py b/pretrain_ict.py deleted file mode 100644 index bf70256bfa1..00000000000 --- a/pretrain_ict.py +++ /dev/null @@ -1,166 +0,0 @@ -# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. - -"""Pretrain BERT for Inverse Cloze Task""" - -from functools import partial -import math - -import torch -import torch.distributed as dist -import torch.nn.functional as F - -from megatron.training import get_args -from megatron.training import print_rank_0 -from megatron.training import get_timers -from megatron.core import mpu -from megatron.core.enums import ModelType -from megatron.legacy.data.biencoder_dataset_utils import get_ict_batch -from megatron.legacy.data.dataset_utils import build_train_valid_test_datasets -from megatron.legacy.model.biencoder_model import biencoder_model_provider -from megatron.training import pretrain -from megatron.training.utils import average_losses_across_data_parallel_group - - -def pretrain_ict_model_provider(pre_process=True, post_process=True): - args = get_args() - - model = biencoder_model_provider( - only_context_model=False, - only_query_model=False, - biencoder_shared_query_context_model=\ - args.biencoder_shared_query_context_model, - pre_process=pre_process, post_process=post_process) - - return model - -def get_group_world_size_rank(): - - group = mpu.get_data_parallel_group() - rank = group.rank() - world_size = group.size() - - return group, rank, world_size - - -class AllgatherFromDataParallelRegion(torch.autograd.Function): - - @staticmethod - def forward(ctx, input_): - assert input_.dim() == 2 - group, rank, world_size = get_group_world_size_rank() - - tensor_list = [torch.empty_like(input_) for _ in range(world_size)] - tensor_list[rank] = input_ - torch.distributed.all_gather(tensor_list, input_, group=group) - - output = torch.cat(tensor_list, dim=0).contiguous() - - return output - - - @staticmethod - def backward(ctx, grad_output): - group, rank, world_size = get_group_world_size_rank() - - assert grad_output.shape[0] % world_size == 0 - dim_size = grad_output.shape[0] // world_size - output_list = torch.split(grad_output, dim_size, dim=0) - - # get chunk from this rank - output = output_list[rank].contiguous() - return output - -def loss_func(output_tensor): - args = get_args() - query_logits, context_logits = output_tensor - - micro_batch_size = query_logits.shape[0] - # recall we assert that tensor_model_parallel_size == 1 - assert mpu.get_tensor_model_parallel_world_size() == 1, \ - "Model parallel size > 1 not supported for ICT" - - global_batch_size = dist.get_world_size() * micro_batch_size - all_query_logits = AllgatherFromDataParallelRegion.apply(query_logits) - all_context_logits = AllgatherFromDataParallelRegion.apply(context_logits) - - # scores are inner products between query and context embeddings - retrieval_scores = torch.matmul(all_query_logits, - torch.transpose(all_context_logits, 0, 1)) - # scaling the retriever scores - if args.retriever_score_scaling: - retrieval_scores = retrieval_scores / math.sqrt(args.hidden_size) - - softmax_scores = F.log_softmax(retrieval_scores, dim=1) - sorted_vals, sorted_indices = torch.topk(softmax_scores, - k=softmax_scores.shape[1], sorted=True) - - def topk_accuracy(k): - return torch.cuda.FloatTensor([sum([int(i in sorted_indices[i, :k]) \ - for i in range(global_batch_size)]) / global_batch_size]) - - topk_accs = [topk_accuracy(int(k)) for k in args.retriever_report_topk_accuracies] - - labels = torch.arange(global_batch_size).long().cuda() - loss = F.nll_loss(softmax_scores, labels, reduction='mean') - reduced_losses = average_losses_across_data_parallel_group([loss, *topk_accs]) - - # Scale the retrieval loss - loss = loss * mpu.get_data_parallel_world_size() - - # create stats_dict with retrieval loss and all specified top-k accuracies - topk_acc_dict = {'top{}_acc'.format(k): v * 100 for k, v in \ - zip(args.retriever_report_topk_accuracies, reduced_losses[1:])} - stats_dict = dict(loss=reduced_losses[0], **topk_acc_dict) - return loss, stats_dict - - - -def forward_step(data_iterator, model): - """Forward step.""" - args = get_args() - timers = get_timers() - - # Get the batch. - timers('batch-generator', log_level=2).start() - query_tokens, query_mask, \ - context_tokens, context_mask, context_indices = get_ict_batch(data_iterator) - timers('batch-generator').stop() - - # Query and Context Types - query_types = torch.cuda.LongTensor(*query_tokens.shape).fill_(0) - context_types = torch.cuda.LongTensor(*context_tokens.shape).fill_(0) - - # Forward model. - output_tensor = model(query_tokens, query_mask, query_types, context_tokens, - context_mask, context_types) - - return output_tensor, partial(loss_func) - -def train_valid_test_datasets_provider(train_val_test_num_samples): - """Build train, valid and test datasets.""" - args = get_args() - print_rank_0('> building train, validation, and test datasets ' - 'for BERT ICT...') - - train_ds, valid_ds, test_ds = build_train_valid_test_datasets( - data_prefix=args.data_path, - splits_string=args.split, - train_valid_test_num_samples=train_val_test_num_samples, - max_seq_length=args.seq_length, - masked_lm_prob=args.mask_prob, - short_seq_prob=args.short_seq_prob, - seed=args.seed, - binary_head=False, - dataset_type='ict') - print_rank_0("> finished creating BERT ICT datasets ...") - - return train_ds, valid_ds, test_ds - - -if __name__ == "__main__": - print_rank_0("WARNING : This script is DEPRECATED. Will be removed in mcore release 0.9") - pretrain(train_valid_test_datasets_provider, - pretrain_ict_model_provider, - ModelType.encoder_or_decoder, - forward_step, - args_defaults={'tokenizer_type': 'BertWordPieceLowerCase'}) diff --git a/pretrain_vision_classify.py b/pretrain_vision_classify.py deleted file mode 100644 index 8d9b28baeb9..00000000000 --- a/pretrain_vision_classify.py +++ /dev/null @@ -1,105 +0,0 @@ -# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. - -"""Pretrain VIT""" - -import torch -import torch.nn.functional as F -from functools import partial -from megatron.training import get_args, get_timers, print_rank_0 -from megatron.core.enums import ModelType -from megatron.legacy.data.vit_dataset import build_train_valid_datasets -from megatron.legacy.model.vision.classification import VitClassificationModel -from megatron.legacy.model.vision.classification import MitClassificationModel -from megatron.training import pretrain -from megatron.training.utils import average_losses_across_data_parallel_group -from megatron.training.arguments import core_transformer_config_from_args - - -def model_provider(pre_process=True, post_process=True): - """Build the model.""" - - args = get_args() - config = core_transformer_config_from_args(args) - if args.vision_backbone_type == 'vit': - print_rank_0("building VIT model ...") - model = VitClassificationModel(config=config, - num_classes=args.num_classes, - pre_process=pre_process, - post_process=post_process) - elif args.vision_backbone_type == 'mit': - print_rank_0("building MIT model ...") - model = MitClassificationModel(num_classes=args.num_classes, - pre_process=pre_process, - post_process=post_process) - else: - raise Exception('{} vision backbone is not supported.'.format( - args.vision_backbone_type)) - return model - - -def get_batch(data_iterator): - """Build the batch.""" - data = next(data_iterator) - - # only data parallelism; no need for broadcast - images = data[0].cuda() - labels = data[1].cuda() - - return images, labels - - -def loss_func(labels, output_tensor): - logits = output_tensor.contiguous().float() - loss = F.cross_entropy(logits, labels) - - outputs = torch.argmax(logits, -1) - correct = (outputs == labels).float() - accuracy = torch.mean(correct) - - averaged_loss = average_losses_across_data_parallel_group([loss, accuracy]) - - return loss, {"loss": averaged_loss[0], "accuracy": averaged_loss[1]} - - -def forward_step(data_iterator, model): - """Forward step.""" - timers = get_timers() - - # Get the batch. - timers("batch-generator", log_level=2).start() - ( - images, - labels, - ) = get_batch(data_iterator) - timers("batch-generator").stop() - - # Forward model. lm_labels - output_tensor = model(images) - - return output_tensor, partial(loss_func, labels) - -def train_valid_test_datasets_provider(train_val_test_num_samples): - """Build train, valid, and test datasets.""" - args = get_args() - - print_rank_0( - "> building train, validation, and test datasets " "for VIT ..." - ) - train_ds, valid_ds = build_train_valid_datasets( - data_path=args.data_path, - image_size=(args.img_h, args.img_w) - ) - print_rank_0("> finished creating VIT datasets ...") - - return train_ds, valid_ds, None - - -if __name__ == "__main__": - - pretrain( - train_valid_test_datasets_provider, - model_provider, - ModelType.encoder_or_decoder, - forward_step, - args_defaults={'dataloader_type': 'cyclic', 'vision_pretraining': True} - ) diff --git a/pretrain_vision_dino.py b/pretrain_vision_dino.py deleted file mode 100644 index f75280c42d7..00000000000 --- a/pretrain_vision_dino.py +++ /dev/null @@ -1,105 +0,0 @@ -# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. - -import torch -import torch.nn.functional as F -import torch.nn as nn -import numpy as np -import torch.distributed as dist -from functools import partial -from megatron.training import get_args, get_timers, print_rank_0 -from megatron.core.enums import ModelType -from megatron.legacy.data.vit_dataset import build_train_valid_datasets -from megatron.legacy.model.vision.dino import DINOPretrainModel -from megatron.legacy.model.vision.knn_monitor import knn_predict, get_feature_bank -from megatron.training import pretrain -from megatron.training.utils import average_losses_across_data_parallel_group, unwrap_model -from megatron.training.arguments import core_transformer_config_from_args - -def model_provider(pre_process=True, post_process=True): - """Build the model.""" - config = core_transformer_config_from_args(get_args()) - return DINOPretrainModel(config, pre_process=pre_process, post_process=post_process) - -def get_batch(data_iterator): - """Build the batch.""" - data = next(data_iterator) - - # only data parallelism; no need for broadcast - if isinstance(data[0], list): - images = [aug.cuda() for aug in data[0]] - else: - images = data[0].cuda() - labels = data[1].cuda() - - return images, labels - - -def loss_func(model, labels, output_tensor, collect_data=False): - args = get_args() - - model = unwrap_model(model) - if model.training: - student_output, teacher_output = output_tensor - loss = model.dino_loss(student_output, teacher_output, args.curr_iteration) - averaged_loss = average_losses_across_data_parallel_group([loss]) - return loss, {"loss": averaged_loss[0]} - else: - _, teacher_feature = output_tensor - feature_bank, feature_labels, classes = get_feature_bank() - feature = F.normalize(teacher_feature.float(), dim=1) - - knn_accs = [] - for k in [10, 20, 100, 200]: - pred_labels = knn_predict(feature, feature_bank, - feature_labels, classes, k, 0.07) - knn_acc = (pred_labels[:, 0] == labels).float().mean() - knn_accs.append(knn_acc) - - averaged_loss = average_losses_across_data_parallel_group(knn_accs) - return 0, {"knn_acc_10": averaged_loss[0], - "knn_acc_20": averaged_loss[1], - "knn_acc_100": averaged_loss[2], - "knn_acc_200": averaged_loss[3]} - - -def forward_step(data_iterator, model): - """Forward step.""" - timers = get_timers() - - # Get the batch. - timers("batch-generator", log_level=2).start() - ( - images, - labels, - ) = get_batch(data_iterator) - timers("batch-generator").stop() - - return model(images), partial(loss_func, model, labels) - - -def train_valid_test_datasets_provider(train_val_test_num_samples): - """Build train, valid, and test datasets.""" - args = get_args() - - print_rank_0( - "> building train, validation, and test datasets " "for VIT ..." - ) - train_ds, valid_ds = build_train_valid_datasets( - data_path=args.data_path, - image_size=(args.img_h, args.img_w) - ) - print_rank_0("> finished creating VIT datasets ...") - - return train_ds, valid_ds, None - - -if __name__ == "__main__": - - pretrain( - train_valid_test_datasets_provider, - model_provider, - ModelType.encoder_or_decoder, - forward_step, - args_defaults={'dataloader_type': 'cyclic', 'vision_pretraining': True} - ) - diff --git a/pretrain_vision_inpaint.py b/pretrain_vision_inpaint.py deleted file mode 100644 index 8570baab5b4..00000000000 --- a/pretrain_vision_inpaint.py +++ /dev/null @@ -1,141 +0,0 @@ -# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. - -"""Pretrain VIT""" - -import torch -import torch.nn.functional as F -from functools import partial -from megatron.training import get_args, get_timers, print_rank_0, print_rank_last -from megatron.core.enums import ModelType -from megatron.legacy.data.vit_dataset import build_train_valid_datasets -from megatron.legacy.model.vision.inpainting import VitInpaintingModel -from megatron.legacy.model.vision.inpainting import MitInpaintingModel -from megatron.training import pretrain -from megatron.training.utils import average_losses_across_data_parallel_group -from tasks.vision.segmentation.metrics import SSIM, PSNR -from megatron.training.arguments import core_transformer_config_from_args - -def model_provider(pre_process=True, post_process=True): - """Build the model.""" - args = get_args() - config = core_transformer_config_from_args(args) - if args.vision_backbone_type == 'vit': - model = VitInpaintingModel(config=config, - pre_process=pre_process, - post_process=post_process) - elif args.vision_backbone_type == 'mit': - model = MitInpaintingModel(config=config, - pre_process=pre_process, - post_process=post_process) - else: - raise Exception('{} vision backbone is not supported.'.format( - args.vision_backbone_type)) - return model - - -def get_batch(data_iterator): - """Build the batch.""" - data = next(data_iterator) - - # only data parallelism; no need for broadcast - images = data[0][0].cuda() - masks = data[0][1].cuda() - return images, masks - - -def loss_func(images, masks, masked_images, outputs, non_loss_data=False): - outputs = outputs.contiguous().float() - masks_flip = 1-masks - flip_masked_outputs = outputs.masked_fill(masks_flip.bool(), 0) - flip_masked_images = images.masked_fill(masks_flip.bool(), 0) - - ssim_fun = SSIM() - psnr_fun = PSNR() - - if not non_loss_data: - mask_count = torch.count_nonzero(masks) - loss = F.mse_loss( - flip_masked_outputs, - flip_masked_images.float(), - reduction="sum" - ) - loss = loss/mask_count - ssim = ssim_fun(flip_masked_outputs, flip_masked_images.float()) - psnr = psnr_fun(flip_masked_outputs, flip_masked_images.float()) - - averaged_loss = average_losses_across_data_parallel_group( - [loss, psnr, ssim] - ) - - return loss, {"loss": averaged_loss[0], - "psnr": averaged_loss[1], - 'ssim': averaged_loss[2]} - else: - synth_images = masked_images.float() + flip_masked_outputs - ssim = ssim_fun(synth_images, images.float()) - psnr = psnr_fun(synth_images, images.float()) - return torch.cat((images, masked_images, synth_images), dim=2), ssim, psnr - - -def forward_step(data_iterator, model): - """Forward step.""" - timers = get_timers() - - # Get the batch. - timers("batch-generator", log_level=2).start() - ( - images, - masks, - ) = get_batch(data_iterator) - timers("batch-generator").stop() - - masked_images = images.masked_fill(masks.bool(), 0) - outputs = model(masked_images) - - # Forward mode - return outputs, partial(loss_func, images, masks, masked_images) - - -def process_non_loss_data(data, iteration, writer): - psnr_sum = 0 - ssim_sum = 0 - for (output_tb, ssim, psnr) in data: - output_tb[output_tb < 0] = 0 - output_tb[output_tb > 1] = 1 - writer.add_images("gt-input-output-vald", output_tb, - global_step=iteration, walltime=None, - dataformats='NCHW') - psnr_sum = psnr_sum + psnr.item() - ssim_sum = ssim_sum + ssim.item() - psnr = psnr_sum/len(data) - ssim = ssim_sum/len(data) - writer.add_scalar('PSNR generate value-validation', psnr, iteration) - writer.add_scalar('SSIM generate value-validation', ssim, iteration) - - -def train_valid_test_datasets_provider(train_val_test_num_samples): - """Build train, valid, and test datasets.""" - args = get_args() - - print_rank_0( - "> building train, validation, and test datasets " "for VIT ..." - ) - train_ds, valid_ds = build_train_valid_datasets( - data_path=args.data_path, - image_size=(args.img_h, args.img_w) - ) - print_rank_0("> finished creating VIT datasets ...") - - return train_ds, valid_ds, None - - -if __name__ == "__main__": - - pretrain( - train_valid_test_datasets_provider, - model_provider, - ModelType.encoder_or_decoder, - forward_step, - process_non_loss_data, - args_defaults={'dataloader_type': 'cyclic', 'vision_pretraining': True} - )