Upload folder using huggingface_hub

bd3493c verified 3 months ago

9.7 kB

"""
NB: this script is for illustration purposes only and is not runnable as our
full dataset is not provided as part of the submission, due to size constraints.
Many of the relevant paths to the data have been thus replaced by dummy paths.
"""

import os
import sys
import pickle
import argparse

import torch
import torch.multiprocessing as mp
from torch.utils.data import DataLoader
import torch.utils.data.distributed
from torch.utils.data.distributed import DistributedSampler
from torch.distributed import init_process_group, destroy_process_group


from trainer import DDPTrainer
from loss_functions import WeightedRmseLoss, PressureWeightedRmseLoss, RmseLoss
from misc_downscaling_functionality import ConvCNPWeatherOnToOff, DownscalingRmseLoss
from loader import *
from models import *
from unet_wrap_padding import *


sys.path.append("../npw/data")
torch.set_float32_matmul_precision("medium")


def ddp_setup(rank, world_size, master_port):
    """
    Args:
        rank: Unique identifier of each process
        world_size: Total number of processes
    """

    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = master_port
    init_process_group(backend="nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)


def start_date(name):
    if name == "train":
        return "2007-01-02"
    elif name == "val":
        return "2019-01-01"
    elif name == "test":
        return "2018-01-01"
    else:
        raise Exception(f"Unrecognised split name {name}")


def end_date(name):
    if name == "train":
        return "2017-12-31"
    elif name == "val":
        return "2019-11-01"
    elif name == "test":
        return "2018-12-21"
    else:
        raise Exception(f"Unrecognised split name {name}")


def main(rank, world_size, output_dir, args):
    """
    Primary training script for the encoder, processor and decoder modules.
    """

    master_port = args.master_port
    lead_time = args.lead_time
    era5_mode = args.era5_mode
    weights_dir = args.weights_dir
    ddp_setup(rank, world_size, master_port)

    # Instantiate loss function
    if args.loss == "lw_rmse":
        lf = WeightedRmseLoss(
            args.res,
            start_ind=args.start_ind,
            end_ind=args.end_ind,
            weight_per_variable=bool(args.weight_per_variable),
        )
    elif args.loss == "lw_rmse_pressure_weighted":
        lf = PressureWeightedRmseLoss(args.res, era5_mode)
    elif args.loss == "rmse":
        lf = RmseLoss()
    elif args.loss == "downscaling_rmse":
        lf = DownscalingRmseLoss()

    # Setup datasets

    # Case 1: training encoder
    if args.mode == "assimilation":
        train_dataset = WeatherDatasetAssimilation(
            device="cuda",
            hadisd_mode="train",
            start_date="2007-01-02",
            end_date="2017-12-31",
            lead_time=0,
            era5_mode="4u",
            res=args.res,
            var_start=args.start_ind,
            var_end=args.end_ind,
            diff=bool(args.diff),
            two_frames=bool(args.two_frames),
        )
        val_dataset = WeatherDatasetAssimilation(
            device="cuda",
            hadisd_mode="train",
            start_date="2019-01-01",
            end_date="2019-12-31",
            lead_time=0,
            era5_mode="4u",
            res=args.res,
            var_start=args.start_ind,
            var_end=args.end_ind,
            diff=bool(args.diff),
            two_frames=bool(args.two_frames),
        )

    # Case 2: training processor
    elif args.mode == "forecast":
        if args.ic == "aardvark":
            train_dataset = FineTuneForecastLoaderNew(
                device="cuda",
                mode="train",
                lead_time=lead_time,
                era5_mode=era5_mode,
                res=args.res,
                frequency=args.frequency,
                diff=bool(args.diff),
                aardvark_ic_path=args.aardvark_ic_path,
                random_lt=True,
            )
            val_dataset = FineTuneForecastLoaderNew(
                device="cuda",
                mode="val",
                lead_time=lead_time,
                era5_mode=era5_mode,
                res=args.res,
                frequency=args.frequency,
                diff=bool(args.diff),
                aardvark_ic_path=args.aardvark_ic_path,
            )
        else:
            train_dataset = ForecastLoader(
                device="cuda",
                mode="train",
                lead_time=lead_time,
                era5_mode=era5_mode,
                res=args.res,
                frequency=args.frequency,
                diff=bool(args.diff),
                u_only=False,
                random_lt=False,
            )
            val_dataset = ForecastLoader(
                device="cuda",
                mode="val",
                lead_time=lead_time,
                era5_mode=era5_mode,
                res=args.res,
                frequency=args.frequency,
                diff=bool(args.diff),
                u_only=False,
                random_lt=False,
            )

    # Case 3: training decoder
    elif args.mode == "downscaling":

        train_dataset = ForecasterDatasetDownscaling(
            start_date="2007-01-02",
            end_date="2017-12-31",
            lead_time=args.lead_time,
            hadisd_var=args.var,
            mode="train",
            device="cuda",
            forecast_path=None,
        )

        val_dataset = ForecasterDatasetDownscaling(
            start_date="2019-01-01",
            end_date="2019-12-21",
            lead_time=args.lead_time,
            hadisd_var=args.var,
            mode="train",
            device="cuda",
            forecast_path=None,
        )

        try:
            os.mkdir(f"{output_dir}lt_{args.lead_time}")
        except FileExistsError:
            pass

        output_dir = f"{output_dir}lt_{args.lead_time}/"

    # Instantiate model

    if args.mode == "downscaling":
        model = ConvCNPWeatherOnToOff(
            in_channels=args.in_channels,
            out_channels=args.end_ind - args.start_ind,
            int_channels=args.int_channels,
            device="cuda",
            res=args.res,
            decoder=args.decoder,
            mode=args.mode,
            film=bool(args.film),
        )
    else:
        model = ConvCNPWeather(
            in_channels=args.in_channels,
            out_channels=args.end_ind - args.start_ind,
            int_channels=args.int_channels,
            device="cuda",
            res=args.res,
            gnp=bool(0),
            decoder=args.decoder,
            mode=args.mode,
            film=bool(args.film),
            two_frames=bool(args.two_frames),
        )

    # Instantiate loaders
    train_sampler = DistributedSampler(train_dataset)
    val_sampler = DistributedSampler(val_dataset)

    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        sampler=train_sampler,
    )

    val_loader = DataLoader(
        val_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        sampler=val_sampler,
    )

    # Instantiate trainer

    trainer = DDPTrainer(
        model,
        rank,
        train_loader,
        val_loader,
        lf,
        output_dir,
        args.lr,
        train_sampler,
        weight_decay=args.weight_decay,
        weights_path=weights_dir,
        tune_film=args.film,
    )

    # Train model

    trainer.train(n_epochs=args.epoch)
    destroy_process_group()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--output_dir")
    parser.add_argument("--mode")
    parser.add_argument("--weights_dir")
    parser.add_argument("--in_channels", type=int)
    parser.add_argument("--out_channels", type=int)
    parser.add_argument("--int_channels", type=int)
    parser.add_argument("--loss")
    parser.add_argument("--ic")
    parser.add_argument("--decoder")
    parser.add_argument("--film")
    parser.add_argument("--aardvark_ic_path")
    parser.add_argument("--two_frames", type=int, default=0)
    parser.add_argument("--weight_per_variable", type=int, default=0)
    parser.add_argument("--batch_size", type=int, default=128)
    parser.add_argument("--epoch", type=int, default=50)
    parser.add_argument("--master_port", default="12345")
    parser.add_argument("--lr", type=float, default=5e-4)
    parser.add_argument("--lead_time", type=int)
    parser.add_argument("--era5_mode", default="4u")
    parser.add_argument("--weight_decay", type=float, default=1e-6)
    parser.add_argument("--res", type=int, default=1)
    parser.add_argument("--frequency", type=int, default=6)
    parser.add_argument("--diff", type=int, default=1)
    parser.add_argument("--start_ind", type=int, default=0)
    parser.add_argument("--end_ind", type=int, default=24)
    parser.add_argument("--downscaling_train_start_date", default="1979-01-01")
    parser.add_argument("--downscaling_train_end_date", default="2017-12-31")
    parser.add_argument("--downscaling_context", default="era5")
    parser.add_argument("--downscaling_lead_time", type=int)
    parser.add_argument("--var", default=None)
    args = parser.parse_args()

    torch.device("cuda")

    # Create results directory
    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.mkdir(output_dir)

    # Save config
    with open(output_dir + "/config.pkl", "wb") as f:
        pickle.dump(vars(args), f)

    world_size = torch.cuda.device_count()
    mp.spawn(main, args=[world_size, output_dir, args], nprocs=world_size)

	"""
	NB: this script is for illustration purposes only and is not runnable as our
	full dataset is not provided as part of the submission, due to size constraints.
	Many of the relevant paths to the data have been thus replaced by dummy paths.
	"""

	import os
	import sys
	import pickle
	import argparse

	import torch
	import torch.multiprocessing as mp
	from torch.utils.data import DataLoader
	import torch.utils.data.distributed
	from torch.utils.data.distributed import DistributedSampler
	from torch.distributed import init_process_group, destroy_process_group


	from trainer import DDPTrainer
	from loss_functions import WeightedRmseLoss, PressureWeightedRmseLoss, RmseLoss
	from misc_downscaling_functionality import ConvCNPWeatherOnToOff, DownscalingRmseLoss
	from loader import *
	from models import *
	from unet_wrap_padding import *


	sys.path.append("../npw/data")
	torch.set_float32_matmul_precision("medium")


	def ddp_setup(rank, world_size, master_port):
	"""
	Args:
	rank: Unique identifier of each process
	world_size: Total number of processes
	"""

	os.environ["MASTER_ADDR"] = "localhost"
	os.environ["MASTER_PORT"] = master_port
	init_process_group(backend="nccl", rank=rank, world_size=world_size)
	torch.cuda.set_device(rank)


	def start_date(name):
	if name == "train":
	return "2007-01-02"
	elif name == "val":
	return "2019-01-01"
	elif name == "test":
	return "2018-01-01"
	else:
	raise Exception(f"Unrecognised split name {name}")


	def end_date(name):
	if name == "train":
	return "2017-12-31"
	elif name == "val":
	return "2019-11-01"
	elif name == "test":
	return "2018-12-21"
	else:
	raise Exception(f"Unrecognised split name {name}")


	def main(rank, world_size, output_dir, args):
	"""
	Primary training script for the encoder, processor and decoder modules.
	"""

	master_port = args.master_port
	lead_time = args.lead_time
	era5_mode = args.era5_mode
	weights_dir = args.weights_dir
	ddp_setup(rank, world_size, master_port)

	# Instantiate loss function
	if args.loss == "lw_rmse":
	lf = WeightedRmseLoss(
	args.res,
	start_ind=args.start_ind,
	end_ind=args.end_ind,
	weight_per_variable=bool(args.weight_per_variable),
	)
	elif args.loss == "lw_rmse_pressure_weighted":
	lf = PressureWeightedRmseLoss(args.res, era5_mode)
	elif args.loss == "rmse":
	lf = RmseLoss()
	elif args.loss == "downscaling_rmse":
	lf = DownscalingRmseLoss()

	# Setup datasets

	# Case 1: training encoder
	if args.mode == "assimilation":
	train_dataset = WeatherDatasetAssimilation(
	device="cuda",
	hadisd_mode="train",
	start_date="2007-01-02",
	end_date="2017-12-31",
	lead_time=0,
	era5_mode="4u",
	res=args.res,
	var_start=args.start_ind,
	var_end=args.end_ind,
	diff=bool(args.diff),
	two_frames=bool(args.two_frames),
	)
	val_dataset = WeatherDatasetAssimilation(
	device="cuda",
	hadisd_mode="train",
	start_date="2019-01-01",
	end_date="2019-12-31",
	lead_time=0,
	era5_mode="4u",
	res=args.res,
	var_start=args.start_ind,
	var_end=args.end_ind,
	diff=bool(args.diff),
	two_frames=bool(args.two_frames),
	)

	# Case 2: training processor
	elif args.mode == "forecast":
	if args.ic == "aardvark":
	train_dataset = FineTuneForecastLoaderNew(
	device="cuda",
	mode="train",
	lead_time=lead_time,
	era5_mode=era5_mode,
	res=args.res,
	frequency=args.frequency,
	diff=bool(args.diff),
	aardvark_ic_path=args.aardvark_ic_path,
	random_lt=True,
	)
	val_dataset = FineTuneForecastLoaderNew(
	device="cuda",
	mode="val",
	lead_time=lead_time,
	era5_mode=era5_mode,
	res=args.res,
	frequency=args.frequency,
	diff=bool(args.diff),
	aardvark_ic_path=args.aardvark_ic_path,
	)
	else:
	train_dataset = ForecastLoader(
	device="cuda",
	mode="train",
	lead_time=lead_time,
	era5_mode=era5_mode,
	res=args.res,
	frequency=args.frequency,
	diff=bool(args.diff),
	u_only=False,
	random_lt=False,
	)
	val_dataset = ForecastLoader(
	device="cuda",
	mode="val",
	lead_time=lead_time,
	era5_mode=era5_mode,
	res=args.res,
	frequency=args.frequency,
	diff=bool(args.diff),
	u_only=False,
	random_lt=False,
	)

	# Case 3: training decoder
	elif args.mode == "downscaling":

	train_dataset = ForecasterDatasetDownscaling(
	start_date="2007-01-02",
	end_date="2017-12-31",
	lead_time=args.lead_time,
	hadisd_var=args.var,
	mode="train",
	device="cuda",
	forecast_path=None,
	)

	val_dataset = ForecasterDatasetDownscaling(
	start_date="2019-01-01",
	end_date="2019-12-21",
	lead_time=args.lead_time,
	hadisd_var=args.var,
	mode="train",
	device="cuda",
	forecast_path=None,
	)

	try:
	os.mkdir(f"{output_dir}lt_{args.lead_time}")
	except FileExistsError:
	pass

	output_dir = f"{output_dir}lt_{args.lead_time}/"

	# Instantiate model

	if args.mode == "downscaling":
	model = ConvCNPWeatherOnToOff(
	in_channels=args.in_channels,
	out_channels=args.end_ind - args.start_ind,
	int_channels=args.int_channels,
	device="cuda",
	res=args.res,
	decoder=args.decoder,
	mode=args.mode,
	film=bool(args.film),
	)
	else:
	model = ConvCNPWeather(
	in_channels=args.in_channels,
	out_channels=args.end_ind - args.start_ind,
	int_channels=args.int_channels,
	device="cuda",
	res=args.res,
	gnp=bool(0),
	decoder=args.decoder,
	mode=args.mode,
	film=bool(args.film),
	two_frames=bool(args.two_frames),
	)

	# Instantiate loaders
	train_sampler = DistributedSampler(train_dataset)
	val_sampler = DistributedSampler(val_dataset)

	train_loader = DataLoader(
	train_dataset,
	batch_size=args.batch_size,
	shuffle=False,
	sampler=train_sampler,
	)

	val_loader = DataLoader(
	val_dataset,
	batch_size=args.batch_size,
	shuffle=False,
	sampler=val_sampler,
	)

	# Instantiate trainer

	trainer = DDPTrainer(
	model,
	rank,
	train_loader,
	val_loader,
	lf,
	output_dir,
	args.lr,
	train_sampler,
	weight_decay=args.weight_decay,
	weights_path=weights_dir,
	tune_film=args.film,
	)

	# Train model

	trainer.train(n_epochs=args.epoch)
	destroy_process_group()


	if __name__ == "__main__":
	parser = argparse.ArgumentParser()
	parser.add_argument("--output_dir")
	parser.add_argument("--mode")
	parser.add_argument("--weights_dir")
	parser.add_argument("--in_channels", type=int)
	parser.add_argument("--out_channels", type=int)
	parser.add_argument("--int_channels", type=int)
	parser.add_argument("--loss")
	parser.add_argument("--ic")
	parser.add_argument("--decoder")
	parser.add_argument("--film")
	parser.add_argument("--aardvark_ic_path")
	parser.add_argument("--two_frames", type=int, default=0)
	parser.add_argument("--weight_per_variable", type=int, default=0)
	parser.add_argument("--batch_size", type=int, default=128)
	parser.add_argument("--epoch", type=int, default=50)
	parser.add_argument("--master_port", default="12345")
	parser.add_argument("--lr", type=float, default=5e-4)
	parser.add_argument("--lead_time", type=int)
	parser.add_argument("--era5_mode", default="4u")
	parser.add_argument("--weight_decay", type=float, default=1e-6)
	parser.add_argument("--res", type=int, default=1)
	parser.add_argument("--frequency", type=int, default=6)
	parser.add_argument("--diff", type=int, default=1)
	parser.add_argument("--start_ind", type=int, default=0)
	parser.add_argument("--end_ind", type=int, default=24)
	parser.add_argument("--downscaling_train_start_date", default="1979-01-01")
	parser.add_argument("--downscaling_train_end_date", default="2017-12-31")
	parser.add_argument("--downscaling_context", default="era5")
	parser.add_argument("--downscaling_lead_time", type=int)
	parser.add_argument("--var", default=None)
	args = parser.parse_args()

	torch.device("cuda")

	# Create results directory
	output_dir = args.output_dir
	if not os.path.exists(output_dir):
	os.mkdir(output_dir)

	# Save config
	with open(output_dir + "/config.pkl", "wb") as f:
	pickle.dump(vars(args), f)

	world_size = torch.cuda.device_count()
	mp.spawn(main, args=[world_size, output_dir, args], nprocs=world_size)