File size: 6,070 Bytes
912fe5a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
import logging
import glob
import torch
import random
import numpy as np
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.distributed import DistributedSampler
from torch import Tensor
import h5py
import math
from my_utils.norm import reshape_fields
def get_data_loader(params, files_pattern, distributed, train):
dataset = GetDataset(params, files_pattern, train)
sampler = DistributedSampler(dataset, shuffle=train) if distributed else None
dataloader = DataLoader(dataset,
batch_size = 1,
num_workers = params.num_data_workers,
shuffle = False,
sampler = sampler if train else None,
drop_last = True,
pin_memory = True)
if train:
return dataloader, dataset, sampler
else:
return dataloader, dataset
class GetDataset(Dataset):
def __init__(self, params, location, train):
self.params = params
self.location = location
self.train = train
self.normalize = params.normalize
self.dt = params.dt
self.n_history = params.n_history
self.in_channels = np.array(params.in_channels)
self.out_channels = np.array(params.out_channels)
self.atmos_channels = np.array(params.atmos_channels)
self.n_in_channels = len(self.in_channels)
self.n_out_channels = len(self.out_channels)
self.add_noise = params.add_noise
self._get_files_stats()
def _get_files_stats(self):
self.files_paths = glob.glob(self.location + "/*.h5")
self.files_paths.sort()
self.n_years = len(self.files_paths)
print('------------', self.files_paths)
with h5py.File(self.files_paths[0], 'r') as _f:
logging.info("Getting file stats from {}".format(self.files_paths[0]))
self.n_samples_per_year = _f['fields'].shape[0] - self.params.multi_steps_finetune
# original image shape (before padding)
self.img_shape_x = _f['fields'].shape[2] - 1 # just get rid of one of the pixels
self.img_shape_y = _f['fields'].shape[3]
self.n_samples_total = self.n_years * self.n_samples_per_year
self.files = [None for _ in range(self.n_years)]
logging.info("Number of samples per year: {}".format(self.n_samples_per_year))
logging.info("Found data at path {}. Number of examples: {}. Image Shape: {} x {} x {}".format(self.location,
self.n_samples_total,
self.img_shape_x,
self.img_shape_y,
self.n_in_channels))
logging.info("Delta t: {} days".format(1 * self.dt))
logging.info("Including {} days of past history in training at a frequency of {} days".format(
1 * self.dt * self.n_history, 1 * self.dt))
def _open_file(self, year_idx):
_file = h5py.File(self.files_paths[year_idx], 'r')
self.files[year_idx] = _file['fields']
def __len__(self):
return self.n_samples_total
def __getitem__(self, global_idx):
year_idx = int(global_idx / self.n_samples_per_year) # which year
local_idx = int(global_idx % self.n_samples_per_year) # which sample in a year
if self.files[year_idx] is None:
self._open_file(year_idx)
if local_idx < self.dt * self.n_history:
local_idx += self.dt * self.n_history
step = 0 if local_idx >= self.n_samples_per_year - self.dt else self.dt
orog = None
if self.params.multi_steps_finetune == 1:
if local_idx == 1463:
local_idx = 1462
if local_idx == 1464:
local_idx = 1463
inp = reshape_fields(
np.nan_to_num(self.files[year_idx][(local_idx-self.dt*self.n_history):(local_idx+1):self.dt, self.in_channels, :120, :240], nan=0),
'inp',
self.params,
self.train,
self.normalize,
orog,
self.add_noise
)
tar = reshape_fields(
np.nan_to_num(self.files[year_idx][local_idx+step, self.out_channels, :120, :240], nan=0),
'tar',
self.params,
self.train,
self.normalize,
orog
)
elif self.params.multi_steps_finetune > 1:
if local_idx == 1463:
local_idx = 1462
if local_idx == 1464:
local_idx = 1463
inp = reshape_fields(
np.nan_to_num(self.files[year_idx][(local_idx-self.dt*self.n_history):(local_idx+1):self.dt, self.in_channels, :120, :240], nan=0),
'inp',
self.params,
self.train,
self.normalize,
orog,
self.add_noise
)
tar = reshape_fields(
np.nan_to_num(self.files[year_idx][local_idx+step:local_idx+step+self.params.multi_steps_finetune, self.out_channels, :120, :240], nan=0),
'tar',
self.params,
self.train,
self.normalize,
orog
)
return inp, tar |