Пример #1
0
 def __init__(self, input_size, num_layers, kernel_size, in_channels, out_channels, dropout_rate, step):
     super(TemporalBlock, self).__init__()
     
     self.dropout_rate = dropout_rate
     self.conv_layers = nn.ModuleList()
     self.dropout_layers = nn.ModuleList()
     self.input_length = input_size[2]
     self.step = step
     
     kernel_size = Util.generate_list_from(kernel_size)
     #factorized kernel
     temporal_kernel_size = [kernel_size[0], 1, 1]
     self.temporal_padding_value = kernel_size[0] // 2        
     temporal_padding = [self.temporal_padding_value, 0, 0]
     intermed_channels = out_channels
     for i in range(num_layers):
         intermed_channels*=2
         if i == (num_layers-1):
             intermed_channels = out_channels 
         self.conv_layers.append(
             nn.Sequential(
                 nn.Conv3d(in_channels, intermed_channels, kernel_size=temporal_kernel_size, 
                           padding=temporal_padding, bias=False),
                 nn.BatchNorm3d(intermed_channels),
                 nn.LeakyReLU(inplace=True)
             )            
         )
         self.dropout_layers.append(nn.Dropout(dropout_rate))
         in_channels = intermed_channels             
Пример #2
0
    def __init__(self, input_size, num_layers, 
    hidden_dim, kernel_size, device, dropout_rate,
    step=5, *args, **kwargs):
        super(MIM, self).__init__()

        self.filter_size = kernel_size
        self.num_hidden_out = input_size[1]
        self.input_length = input_size[2]
        self.step = step
        self.num_layers = num_layers
        self.num_hidden = Util.generate_list_from(hidden_dim,num_layers)
        self.device = device
        self.stlstm_layer = nn.ModuleList()
        self.stlstm_layer_diff = nn.ModuleList()
        
        num_hidden_in = self.num_hidden_out
        for i in range(self.num_layers):
            if i < 1:
                self.stlstm_layer.append(
                    SpatioTemporalLSTMCell(self.filter_size, num_hidden_in, self.num_hidden[i], input_size, device, dropout_rate)
                )
            else:
                self.stlstm_layer.append(
                    MIMBlock(self.filter_size, self.num_hidden[i], input_size, device, dropout_rate)
                )
                     
        for i in range(self.num_layers - 1):
            self.stlstm_layer_diff.append(
                MIMS(self.filter_size, self.num_hidden[i+1], input_size, device, dropout_rate)
            )

        self.conv_last = nn.Conv2d(self.num_hidden[num_layers - 1], self.num_hidden_out,
                                   kernel_size=1, stride=1, padding=0, bias=False) 
Пример #3
0
    def __init__(self,
                 input_size,
                 num_layers,
                 hidden_dim,
                 kernel_size,
                 device,
                 dropout_rate,
                 step=5):
        super(PredRNN, self).__init__()

        self.frame_channel = input_size[1]
        self.num_layers = num_layers
        self.num_hidden = Util.generate_list_from(hidden_dim, num_layers)
        self.device = device
        cell_list = []
        self.input_length = input_size[2]
        self.step = step
        width = input_size[4]
        in_channel = self.frame_channel
        for i in range(self.num_layers):
            in_channel = self.frame_channel if i == 0 else self.num_hidden[i -
                                                                           1]
            cell_list.append(
                SpatioTemporalLSTMCell(in_channel, self.num_hidden[i], width,
                                       kernel_size, 1, dropout_rate))

        self.cell_list = nn.ModuleList(cell_list)
        self.conv_last = nn.Conv2d(self.num_hidden[num_layers - 1],
                                   self.frame_channel,
                                   kernel_size=1,
                                   stride=1,
                                   padding=0,
                                   bias=False)
Пример #4
0
    def __init__(self, num_layers, kernel_size, in_channels, out_channels,
                 dropout_rate):
        super(SpatialBlock, self).__init__()
        self.padding = kernel_size // 2
        self.dropout_rate = dropout_rate
        self.conv_layers = nn.ModuleList()
        self.dropout_layers = nn.ModuleList()

        kernel_size = Util.generate_list_from(kernel_size)
        #factorized kernel
        spatial_kernel_size = [1, kernel_size[1], kernel_size[2]]
        spatial_padding_value = kernel_size[1] // 2
        spatial_padding = [0, spatial_padding_value, spatial_padding_value]
        intermed_channels = out_channels
        for i in range(num_layers):
            intermed_channels *= 2
            if i == (num_layers - 1):
                intermed_channels = out_channels
            self.conv_layers.append(
                nn.Sequential(
                    nn.Conv3d(in_channels,
                              intermed_channels,
                              kernel_size=spatial_kernel_size,
                              padding=spatial_padding,
                              bias=False), nn.BatchNorm3d(intermed_channels),
                    nn.LeakyReLU(inplace=True)))
            self.dropout_layers.append(nn.Dropout(dropout_rate))
            in_channels = intermed_channels
Пример #5
0
 def __init__(self, input_size, num_layers, kernel_size, in_channels, out_channels, dropout_rate, step):
     super(TemporalCausalBlock_NoChannelIncrease, self).__init__()
     
     self.dropout_rate = dropout_rate
     self.conv_layers = nn.ModuleList()
     self.lrelu_layers = nn.ModuleList()
     self.batch_layers = nn.ModuleList()
     self.dropout_layers = nn.ModuleList()
     self.input_length = input_size[2]
     self.step = step
     
     kernel_size = Util.generate_list_from(kernel_size)
     #factorized kernel
     temporal_kernel_size = [kernel_size[0], 1, 1]
     self.temporal_padding_value = kernel_size[0] - 1      
     temporal_padding = [self.temporal_padding_value, 0, 0]
     for i in range(num_layers):
         self.conv_layers.append(
             nn.Conv3d(in_channels, out_channels, kernel_size=temporal_kernel_size, 
                       padding=temporal_padding, bias=False)
         )
         self.lrelu_layers.append(nn.LeakyReLU())
         self.batch_layers.append(nn.BatchNorm3d(out_channels))
         self.dropout_layers.append(nn.Dropout(dropout_rate))
         in_channels = out_channels
Пример #6
0
    def __init__(self, kernel_size, in_channels, out_channels, dropout_rate,
                 bias):
        super(Conv2Plus1Block, self).__init__()

        kernel_size = Util.generate_list_from(kernel_size)
        #factorized kernel
        spatial_kernel_size = [1, kernel_size[1], kernel_size[2]]
        temporal_kernel_size = [kernel_size[0], 1, 1]

        spatial_padding_value = kernel_size[1] // 2
        temporal_padding_value = kernel_size[0] // 2

        spatial_padding = [0, spatial_padding_value, spatial_padding_value]
        temporal_padding = [temporal_padding_value, 0, 0]

        intermed_channels = int(math.floor((kernel_size[0] * kernel_size[1] * kernel_size[2] * in_channels * out_channels)/ \
                            (kernel_size[1] * kernel_size[2] * in_channels + kernel_size[0] * out_channels)))

        self.spatial_conv = nn.Sequential(
            nn.Conv3d(in_channels,
                      intermed_channels,
                      spatial_kernel_size,
                      padding=spatial_padding,
                      bias=bias), nn.BatchNorm3d(intermed_channels),
            nn.LeakyReLU(inplace=True))
        self.temporal_conv = nn.Conv3d(intermed_channels,
                                       out_channels,
                                       temporal_kernel_size,
                                       padding=temporal_padding,
                                       bias=bias)
 def __init__(self, kernel_size, in_channels, out_channels):
     super().__init__()
     kernel_size = Util.generate_list_from(kernel_size)
     spatial_value = kernel_size[1] // 2
     kernel_size = [1, spatial_value, spatial_value]
     stride = [1, spatial_value, spatial_value]
     self.conv = nn.Sequential(
         nn.ConvTranspose3d(in_channels,
                            out_channels,
                            kernel_size=kernel_size,
                            stride=stride), nn.LeakyReLU(inplace=True))
 def __init__(self, kernel_size, in_channels, out_channels):
     super().__init__()
     kernel_size = Util.generate_list_from(kernel_size)
     temporal_value = kernel_size[0] // 2
     spatial_value = kernel_size[1] // 2
     padding = [temporal_value, spatial_value, spatial_value]
     stride = [1, spatial_value, spatial_value]
     self.down_block = nn.Sequential(
         nn.Conv3d(in_channels,
                   out_channels,
                   kernel_size=kernel_size,
                   padding=padding,
                   stride=stride), nn.LeakyReLU(inplace=True))
Пример #9
0
    def __init__(self, kernel_size, in_channels, out_channels, dropout_rate,
                 bias):
        super(Conv3DBlock, self).__init__()

        kernel_size = Util.generate_list_from(kernel_size)
        spatial_padding_value = kernel_size[1] // 2
        temporal_padding_value = kernel_size[0] // 2

        padding = [
            temporal_padding_value, spatial_padding_value,
            spatial_padding_value
        ]

        self.conv = nn.Conv3d(in_channels,
                              out_channels,
                              kernel_size,
                              padding=padding,
                              bias=bias)
Пример #10
0
    def __init__(self, input_size, kernel_size, in_channels, out_channels,
                 dropout_rate, step):
        super(TemporalGeneratorBlock, self).__init__()
        self.step = step
        self.input_length = input_size[2]
        self.tconv_layers = nn.ModuleList()
        self.conv_layers = nn.ModuleList()
        kernel_size = Util.generate_list_from(kernel_size)
        #factorized kernel
        spatial_kernel_size = [1, kernel_size[1], kernel_size[2]]
        spatial_padding_value = kernel_size[1] // 2
        spatial_padding = [0, spatial_padding_value, spatial_padding_value]

        num_layers = math.ceil(
            (self.step - self.input_length) / (2 * self.input_length))
        intermed_channels = out_channels
        for i in range(num_layers):
            intermed_channels *= 2
            if i == (num_layers - 1):
                intermed_channels = out_channels
            self.tconv_layers.append(
                nn.Sequential(
                    nn.ConvTranspose3d(in_channels,
                                       intermed_channels, [4, 1, 1],
                                       stride=[2, 1, 1],
                                       padding=[1, 0, 0],
                                       bias=False),
                    nn.BatchNorm3d(intermed_channels),
                    nn.LeakyReLU(inplace=True)))
            in_channels = intermed_channels

        num_layers = self.step // self.input_length
        intermed_channels *= 2
        for i in range(num_layers):
            self.conv_layers.append(
                nn.Sequential(
                    nn.Conv3d(in_channels,
                              intermed_channels,
                              kernel_size=spatial_kernel_size,
                              padding=spatial_padding,
                              bias=False), nn.BatchNorm3d(intermed_channels),
                    nn.LeakyReLU(inplace=True)))
            in_channels = intermed_channels
            intermed_channels = out_channels
Пример #11
0
 def __init__(self, input_size, num_layers, kernel_size, in_channels, out_channels, dropout_rate, step):
     super(TemporalReversedBlock_NoChannelIncrease, self).__init__()
     
     self.dropout_rate = dropout_rate
     self.conv_layers = nn.ModuleList()
     self.lrelu_layers = nn.ModuleList()
     self.batch_layers = nn.ModuleList()
     self.dropout_layers = nn.ModuleList()        
     self.input_length = input_size[2]
     self.step = step
     
     kernel_size = Util.generate_list_from(kernel_size)
     #factorized kernel
     temporal_kernel_size = [kernel_size[0], 1, 1]          
     for i in range(num_layers):
         self.conv_layers.append(
        	    RNet(in_channels, out_channels, kernel_size=temporal_kernel_size, bias=False)
         )
         self.dropout_layers.append(nn.Dropout(dropout_rate))
         in_channels = out_channels
Пример #12
0
    def __execute_learning(self, model, criterion, optimizer, train_loader,
                           val_loader, test_loader, dataset_name,
                           filename_prefix, dropout_rate):
        criterion_name = type(criterion).__name__
        filename_prefix += '_' + criterion_name
        util = Util(self.model_descr, self.dataset_type, self.version,
                    filename_prefix)

        # Training the model
        checkpoint_filename = util.get_checkpoint_filename()
        trainer = Trainer(model, criterion, optimizer, train_loader,
                          val_loader, self.epochs, self.device, self.verbose,
                          self.patience, self.no_stop)

        start_timestamp = tm.time()
        train_losses, val_losses = trainer.fit(checkpoint_filename)
        end_timestamp = tm.time()
        # Error analysis
        util.save_loss(train_losses, val_losses)
        util.plot([train_losses, val_losses], ['Training', 'Validation'],
                  'Epochs', 'Error', 'Error analysis', self.plot)

        # Load model with minimal loss after training phase
        model, _, best_epoch, val_loss = trainer.load_checkpoint(
            checkpoint_filename)

        # Evaluating the model
        evaluator = Evaluator(model, criterion, test_loader, self.device)
        test_loss = evaluator.eval()
        train_time = end_timestamp - start_timestamp
        print(
            f'Training time: {util.to_readable_time(train_time)}\n{self.model_descr} {criterion_name}: {test_loss:.4f}\n'
        )

        return {
            'best_epoch': best_epoch,
            'val_error': val_loss,
            'test_error': test_loss,
            'train_time': train_time,
            'loss_type': criterion_name,
            'dropout_rate': dropout_rate,
            'dataset': dataset_name
        }
Пример #13
0

if __name__ == '__main__':
    args = get_arguments()
    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.cuda)
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    message, model_descr = None, None
    if (args.convlstm):
        model_descr = 'ConvLSTM'
    else:
        model_descr = 'STConvS2s'

    model_builder = MLBuilder(model_descr, args.version, args.plot,
                              args.no_seed, args.verbose, args.small_dataset,
                              args.no_stop, args.epoch, args.patience, device,
                              args.workers, args.convlstm, args.mae,
                              args.chirps, args.step)

    print(f'RUN MODEL: {model_descr}')
    print(f'Device: {device}')
    # start time is saved when creating an instance of Util
    util = Util(model_descr, version=args.version)
    try:
        message = run(model_builder, args.iteration, util)
        message['step'] = args.step
        message['hostname'] = platform.node()
    except Exception as e:
        traceback.print_exc()
        message = '=> Error: ' + str(e)
    util.send_email(message, args.email)
Пример #14
0
    def run_model(self, number):
        self.__define_seed(number)
        validation_split = 0.2
        test_split = 0.2
        # Loading the dataset
        ds = xr.open_mfdataset(self.dataset_file)
        if (self.config.small_dataset):
            ds = ds[dict(sample=slice(0, 500))]

        train_dataset = NetCDFDataset(ds,
                                      test_split=test_split,
                                      validation_split=validation_split)
        val_dataset = NetCDFDataset(ds,
                                    test_split=test_split,
                                    validation_split=validation_split,
                                    is_validation=True)
        test_dataset = NetCDFDataset(ds,
                                     test_split=test_split,
                                     validation_split=validation_split,
                                     is_test=True)
        if self.config.verbose:
            print('[X_train] Shape:', train_dataset.X.shape)
            print('[y_train] Shape:', train_dataset.y.shape)
            print('[X_val] Shape:', val_dataset.X.shape)
            print('[y_val] Shape:', val_dataset.y.shape)
            print('[X_test] Shape:', test_dataset.X.shape)
            print('[y_test] Shape:', test_dataset.y.shape)
            print(
                f'Train on {len(train_dataset)} samples, validate on {len(val_dataset)} samples'
            )

        params = {
            'batch_size': self.config.batch,
            'num_workers': self.config.workers,
            'worker_init_fn': self.__init_seed
        }

        train_loader = DataLoader(dataset=train_dataset,
                                  shuffle=True,
                                  **params)
        val_loader = DataLoader(dataset=val_dataset, shuffle=False, **params)
        test_loader = DataLoader(dataset=test_dataset, shuffle=False, **params)

        models = {
            'stconvs2s-r': stconvs2s.STConvS2S_R,
            'stconvs2s-c': stconvs2s.STConvS2S_C,
            'convlstm': stconvlstm.STConvLSTM,
            'predrnn': predrnn.PredRNN,
            'mim': mim.MIM,
            'conv2plus1d': conv2plus1d.Conv2Plus1D,
            'conv3d': conv3d.Conv3D,
            'enc-dec3d': encoder_decoder3d.Endocer_Decoder3D,
            'vlstm': nosocial.VanillaLSTM_Downsample,
            'slstm': slstm.SocialLSTM_Downsample,
            'sclstm': sclstm.SocialConvLSTM,
        }
        if self.config.model not in models:
            raise ValueError(
                f'{self.config.model} is not a valid model name. Choose between: {models.keys()}'
            )
            quit()

        # Creating the model
        model_bulder = models[self.config.model]
        model = model_bulder(
            input_size=train_dataset.X.shape,
            num_layers=self.config.num_layers,
            hidden_dim=self.config.hidden_dim,
            kernel_size=self.config.kernel_size,
            device=self.device,
            dropout_rate=self.dropout_rate,
            step=int(self.step),
            share=self.config.share,
            lstms_shape=self.config.lstms_shape,
        )
        model.to(self.device)
        metrics = {
            'rmse': (RMSELoss, RMSEDownSample),
            'mae': (L1Loss, L1LossDownSample)
        }
        if self.config.loss not in metrics:
            raise ValueError(
                f'{self.config.loss} is not a valid loss function name. Choose between: {models.keys()}'
            )
            quit()
        if self.config.model in ['vlstm', 'slstm']:
            loss = metrics[self.config.loss][1](train_dataset.X.shape)
        else:
            loss = metrics[self.config.loss][0]()
        loss.to(self.device)

        opt_params = {'lr': 0.001, 'alpha': 0.9, 'eps': 1e-6}
        optimizer = torch.optim.RMSprop(model.parameters(), **opt_params)
        util = Util(self.config.model, self.dataset_type, self.config.version,
                    self.filename_prefix)
        train_info = {'train_time': 0}
        if self.config.pre_trained is None:
            train_info = self.__execute_learning(model, loss, optimizer,
                                                 train_loader, val_loader,
                                                 util)

        eval_info = self.__load_and_evaluate(model, loss, optimizer,
                                             test_loader,
                                             train_info['train_time'], util)

        if (torch.cuda.is_available()):
            torch.cuda.empty_cache()

        return {**train_info, **eval_info}
Пример #15
0
    def run_model(self, number):
        self.__define_seed(number)
        validation_split = 0.2
        test_split = 0.2
        # Loading the dataset
        ds = xr.open_mfdataset(self.dataset_file)
        if (self.config.small_dataset):
            ds = ds[dict(sample=slice(0, 500))]

        train_dataset = NetCDFDataset(ds,
                                      test_split=test_split,
                                      validation_split=validation_split,
                                      x_step=self.x_step)
        val_dataset = NetCDFDataset(ds,
                                    test_split=test_split,
                                    validation_split=validation_split,
                                    x_step=self.x_step,
                                    is_validation=True)
        test_dataset = NetCDFDataset(ds,
                                     test_split=test_split,
                                     validation_split=validation_split,
                                     x_step=self.x_step,
                                     is_test=True)

        util = Util(self.config.model, self.dataset_type, self.config.version,
                    self.filename_prefix)

        # normalizing data
        num_channels = train_dataset.X.shape[1]
        if num_channels > 1:
            normalizer_x = Normalizer()
            normalizer_x.observe(train_dataset.X)
            normalizer_y = Normalizer()
            normalizer_y.observe(train_dataset.y)

            train_dataset.X = normalizer_x.normalize(train_dataset.X)
            train_dataset.y = normalizer_y.normalize(train_dataset.y)

            val_dataset.X = normalizer_x.normalize(val_dataset.X)
            val_dataset.y = normalizer_y.normalize(val_dataset.y)

            test_dataset.X = normalizer_x.normalize(test_dataset.X)
            test_dataset.y = normalizer_y.normalize(test_dataset.y)

            util.save_normalization_parameters(normalizer_x, normalizer_y)

        # INITIAL STATE - batch x channel x time x latitude x longitude
        initial_state = torch.tensor(train_dataset.X)[:1, :, :1].to(
            self.device)

        if (self.config.verbose):
            print('[X_train] Shape:', train_dataset.X.shape)
            print('[y_train] Shape:', train_dataset.y.shape)
            print('[X_val] Shape:', val_dataset.X.shape)
            print('[y_val] Shape:', val_dataset.y.shape)
            print('[X_test] Shape:', test_dataset.X.shape)
            print('[y_test] Shape:', test_dataset.y.shape)
            print(
                f'Train on {len(train_dataset)} samples, validate on {len(val_dataset)} samples'
            )

        params = {
            'batch_size': self.config.batch,
            'num_workers': self.config.workers,
            'worker_init_fn': self.__init_seed
        }

        train_loader = DataLoader(dataset=train_dataset,
                                  shuffle=True,
                                  **params)
        val_loader = DataLoader(dataset=val_dataset, shuffle=False, **params)
        test_loader = DataLoader(dataset=test_dataset, shuffle=False, **params)

        models = {
            'stconvs2s-r': STConvS2S_R,
            'stconvs2s-c': STConvS2S_C,
            'convlstm': STConvLSTM,
            'predrnn': PredRNN,
            'mim': MIM,
            'conv2plus1d': Conv2Plus1D,
            'conv3d': Conv3D,
            'enc-dec3d': Endocer_Decoder3D,
            'ablation-stconvs2s-nocausalconstraint':
            AblationSTConvS2S_NoCausalConstraint,
            'ablation-stconvs2s-notemporal': AblationSTConvS2S_NoTemporal,
            'ablation-stconvs2s-r-nochannelincrease':
            AblationSTConvS2S_R_NoChannelIncrease,
            'ablation-stconvs2s-c-nochannelincrease':
            AblationSTConvS2S_C_NoChannelIncrease,
            'ablation-stconvs2s-r-inverted': AblationSTConvS2S_R_Inverted,
            'ablation-stconvs2s-c-inverted': AblationSTConvS2S_C_Inverted,
            'ablation-stconvs2s-r-notfactorized':
            AblationSTConvS2S_R_NotFactorized,
            'ablation-stconvs2s-c-notfactorized':
            AblationSTConvS2S_C_NotFactorized
        }
        if not (self.config.model in models):
            raise ValueError(
                f'{self.config.model} is not a valid model name. Choose between: {models.keys()}'
            )
            quit()

        # Creating the model
        model_bulder = models[self.config.model]
        model = model_bulder(train_dataset.X.shape, self.config.num_layers,
                             self.config.hidden_dim, self.config.kernel_size,
                             self.device, self.dropout_rate, self.y_step)

        # Use all disponible GPUs
        if torch.cuda.device_count() > 1:
            print("Let's use", torch.cuda.device_count(), "GPUs!")
            model = nn.DataParallel(model)

        model.to(self.device)
        criterion = RMSELoss(reg=self.config.regularization,
                             initial_state=initial_state)
        opt_params = {
            'lr': self.config.learning_rate,
            'alpha': 0.9,
            'eps': 1e-6
        }
        optimizer = torch.optim.RMSprop(model.parameters(), **opt_params)

        train_info = {'train_time': 0}
        if self.config.pre_trained is None:
            train_info = self.__execute_learning(model, criterion, optimizer,
                                                 train_loader, val_loader,
                                                 util)

        eval_info = self.__load_and_evaluate(model, criterion, optimizer,
                                             test_loader,
                                             train_info['train_time'], util)

        if (torch.cuda.is_available()):
            torch.cuda.empty_cache()

        return {**train_info, **eval_info}
Пример #16
0

if __name__ == '__main__':
    print('RUN MODEL: ARIMA')
    args = get_arguments()
    dataset_name, dataset_file = get_dataset_file(args.chirps)
    ds = xr.open_mfdataset(dataset_file)
    with Pool() as pool:
        i = range(ds.lat.size)
        index_list = list(itertools.product(i, i))
        # separate time series based on each location
        ds_list = [
            ds.isel(lat=index[0], lon=index[1]).to_dataframe()
            for index in index_list
        ]
        util = Util('ARIMA')
        results = pool.starmap(
            run_arima,
            zip(ds_list, itertools.repeat(args.chirps),
                itertools.repeat(int(args.step))))
        results = np.array(results)
        pool.close()
        pool.join()

        print('Elapsed time', util.get_time_info()['elapsed_time'])
        rmse_list = [result[0] for result in results if result[0] >= 0]
        mae_list = [result[1] for result in results if result[1] >= 0]
        rmse_mean, rmse_std = np.mean(rmse_list), np.std(rmse_list)
        mae_mean, mae_std = np.mean(mae_list), np.std(mae_list)
        print('\nRMSE: ', rmse_list)
        print('\nMAE: ', mae_list)
Пример #17
0
    def run_model(self, number):
        self.__define_seed(number)
        validation_split = 0.2
        test_split = 0.2
        # Loading the dataset
        ds = xr.open_mfdataset(self.dataset_file)
        if (self.config.small_dataset):
            ds = ds[dict(sample=slice(0, 500))]

        train_dataset = NetCDFDataset(ds,
                                      test_split=test_split,
                                      validation_split=validation_split)
        val_dataset = NetCDFDataset(ds,
                                    test_split=test_split,
                                    validation_split=validation_split,
                                    is_validation=True)
        test_dataset = NetCDFDataset(ds,
                                     test_split=test_split,
                                     validation_split=validation_split,
                                     is_test=True)
        if (self.config.verbose):
            print('[X_train] Shape:', train_dataset.X.shape)
            print('[y_train] Shape:', train_dataset.y.shape)
            print('[X_val] Shape:', val_dataset.X.shape)
            print('[y_val] Shape:', val_dataset.y.shape)
            print('[X_test] Shape:', test_dataset.X.shape)
            print('[y_test] Shape:', test_dataset.y.shape)
            print(
                f'Train on {len(train_dataset)} samples, validate on {len(val_dataset)} samples'
            )

        params = {
            'batch_size': self.config.batch,
            'num_workers': self.config.workers,
            'worker_init_fn': self.__init_seed
        }

        train_loader = DataLoader(dataset=train_dataset,
                                  shuffle=True,
                                  **params)
        val_loader = DataLoader(dataset=val_dataset, shuffle=False, **params)
        test_loader = DataLoader(dataset=test_dataset, shuffle=False, **params)

        models = {
            'stconvs2s-r': STConvS2S_R,
            'stconvs2s-c': STConvS2S_C,
            'convlstm': STConvLSTM,
            'predrnn': PredRNN,
            'mim': MIM,
            'conv2plus1d': Conv2Plus1D,
            'conv3d': Conv3D,
            'enc-dec3d': Endocer_Decoder3D,
            'ablation-stconvs2s-nocausalconstraint':
            AblationSTConvS2S_NoCausalConstraint,
            'ablation-stconvs2s-notemporal': AblationSTConvS2S_NoTemporal,
            'ablation-stconvs2s-r-nochannelincrease':
            AblationSTConvS2S_R_NoChannelIncrease,
            'ablation-stconvs2s-c-nochannelincrease':
            AblationSTConvS2S_C_NoChannelIncrease,
            'ablation-stconvs2s-r-inverted': AblationSTConvS2S_R_Inverted,
            'ablation-stconvs2s-c-inverted': AblationSTConvS2S_C_Inverted,
            'ablation-stconvs2s-r-notfactorized':
            AblationSTConvS2S_R_NotFactorized,
            'ablation-stconvs2s-c-notfactorized':
            AblationSTConvS2S_C_NotFactorized
        }
        if not (self.config.model in models):
            raise ValueError(
                f'{self.config.model} is not a valid model name. Choose between: {models.keys()}'
            )
            quit()

        # Creating the model
        model_bulder = models[self.config.model]
        model = model_bulder(train_dataset.X.shape, self.config.num_layers,
                             self.config.hidden_dim, self.config.kernel_size,
                             self.device, self.dropout_rate, int(self.step))
        model.to(self.device)
        criterion = RMSELoss()
        opt_params = {'lr': 0.001, 'alpha': 0.9, 'eps': 1e-6}
        optimizer = torch.optim.RMSprop(model.parameters(), **opt_params)
        util = Util(self.config.model, self.dataset_type, self.config.version,
                    self.filename_prefix)

        train_info = {'train_time': 0}
        if self.config.pre_trained is None:
            train_info = self.__execute_learning(model, criterion, optimizer,
                                                 train_loader, val_loader,
                                                 util)

        eval_info = self.__load_and_evaluate(model, criterion, optimizer,
                                             test_loader,
                                             train_info['train_time'], util)

        if (torch.cuda.is_available()):
            torch.cuda.empty_cache()

        return {**train_info, **eval_info}
Пример #18
0
                                  train_times_epochs, iteration, util)
    new_model_info['dataset'] = model_info['dataset']
    return new_model_info


if __name__ == '__main__':
    args = get_arguments()
    #os.environ["CUDA_VISIBLE_DEVICES"]=str(args.cuda)
    device = torch.device('cpu')
    device_descr = 'CPU'
    if torch.cuda.is_available():
        device = torch.device('cuda')
        device_descr = 'GPU'

    message = None
    model_builder = MLBuilder(args, device)
    print(f'RUN MODEL: {args.model.upper()}')
    print(f'Device: {device_descr}')
    print(f'Settings: {args}')
    # start time is saved when creating an instance of Util
    util = Util(args.model, version=args.version)
    try:
        message = run(model_builder, args.iteration, util)
        message['x_step'] = args.x_step
        message['y_step'] = args.y_step
        message['hostname'] = platform.node()
    except Exception as e:
        traceback.print_exc()
        message = '=> Error: ' + str(e)
    util.send_email(message, args.email)