Function fast_learning

learner.py:83–121

(law_list, model, model_name, p_epoch, bs, train_occupancy, train_price, seq_l, pre_l, device, adj_dense)

def fast_learning(law_list, model, model_name, p_epoch, bs, train_occupancy, train_price, seq_l, pre_l, device, adj_dense):
    n_laws = len(law_list)
    fast_datasets = dict()
    fast_loaders = dict()
    for n in range(n_laws):
        fast_datasets[n] = fn.CreateFastDataset(train_occupancy, train_price, seq_l, pre_l, law_list[n], device, adj_dense)
        fast_loaders[n] = DataLoader(fast_datasets[n], batch_size=bs, shuffle=True, drop_last=True)
    
    optimizer = torch.optim.Adam(model.parameters(), weight_decay=0.00001)
    loss_function = torch.nn.MSELoss()
    for epoch in tqdm(range(p_epoch), desc='Pre-training'):
        for n in range(n_laws):
            for j, data in enumerate(fast_loaders[n]):
                '''
                occupancy = (batch, seq, node)
                price = (batch, seq, node)
                label = (batch, node)
                '''
                occupancy, price, label, prc_ch, label_ch = data
                optimizer.zero_grad()
                predict = model(occupancy, prc_ch)
                loss = loss_function(predict, label_ch)
                loss.backward()
                optimizer.step()

            for j, data in enumerate(fast_loaders[n]):
                '''
                occupancy = (batch, seq, node)
                price = (batch, seq, node)
                label = (batch, node)
                '''
                occupancy, price, label, prc_ch, label_ch = data
                optimizer.zero_grad()
                predict = model(occupancy, prc_ch)
                loss = loss_function(predict, label_ch)
                loss.backward()
                optimizer.step()

    return model