driver.py

from hour_dataset import HourDataset
from lstm import LSTMSimple
import torch
from torch import dropout, nn
import numpy as np
import sys
from fitlog import FitLog
from torch.utils.data import DataLoader
import os
import torch.distributed
from sklearn.metrics import accuracy_score
# from metric import MetricCollector


class LSTMDriverHour():

    def __init__(self, local_rank):
        # self.epochs = 20000
        self.epochs = 2000
        self.batch_size = 156
        self.hidden_size = 2048
        self.nlayer = 2
        self.n_class = 2
        self.lr = 0.0001
        self.early_stop_thresh = 0.001
        self.early_stop_nreach_limit = 10

        self.dataset = None
        self.model = None
        self.criterion = None
        self.optimizer = None
        self.loader = None
        self.device = None

        self.local_rank = local_rank

    def __get_description(self):
        ret = {'epochs': self.epochs, \
               'batch_size:': self.batch_size, \
               'hidden_size': self.hidden_size, \
               'nlayer': self.nlayer, \
               'n_class': self.n_class, \
               'lr': self.lr, \
               'step_size': self.dataset.get_time_step(), \
               'input_size': self.dataset.get_input_size()}
        return ret

    def load(self, path=None):
        self.dataset = HourDataset(100, random_seed=1)
        if path == None:
            self.dataset.load()
        else:
            self.dataset.loadfile(path)

    def __prepare_for_fit(self):
        self.device = torch.device('cuda', self.local_rank)
        #self.device = torch.device('cuda', 0)

        self.model = LSTMSimple(self.dataset.get_input_size(), self.hidden_size, self.nlayer, self.n_class, self.device)

        self.model.to(self.device)
        self.model = nn.parallel.DistributedDataParallel(self.model, device_ids=[self.local_rank],find_unused_parameters=True)

        self.criterion = nn.CrossEntropyLoss()
        self.optimizer = torch.optim.Adam(self.model.parameters(), self.lr)
        # self.optimizer = torch.optim.SGD(self.model.parameters(), self.lr)
    
    def get_device(self):
        device_mlu = None
        device_gpu = None
        try:
            # device_mlu = torch.device('mlu')
            device_gpu = torch.device('cuda')
        except Exception as err:
            print(err)
        if device_mlu:
            self.fitlog.append('mlu', True, True)
            return device_mlu
        elif device_gpu:
            self.fitlog.append('cuda', True, True)
            return device_gpu
        else:
            self.fitlog.append('cpu', True, True)
            return torch.device('cpu')

    def fit(self):
        if self.dataset == None:
            print('no data was loaded')
            return None

        self.__prepare_for_fit()
        if self.local_rank == 0:
            fitlog = FitLog(folderpath='logs/')
            fitlog.append(str(self.__get_description()), with_time=True)
        print(self.__get_description())

        self.dataset.set_mode('train')

        sampler = torch.utils.data.distributed.DistributedSampler(self.dataset)
        self.loader = torch.utils.data.DataLoader(self.dataset, batch_size=self.batch_size, sampler=sampler,
                                                  shuffle=False)
        # self.loader = DataLoader(self.dataset, batch_size=self.batch_size, shuffle=True)
        nbatch = len(self.loader)
        print('batch:', nbatch)

        lowest_loss = 99999999
        lowest_loss_at = -1
        early_stop_nreach = 0
        best_acc = 0
        best_acc_at = -1
        import time


        dt = [0.0, 0.0, 0.0]
        ts = time.time()
        # 如果是训练单个epoch，也可以移到下一层循环内
        best_acc=0
        for epoch in range(self.epochs):
            self.dataset.set_mode('train')
            self.model.train()
            t0 = time.time()
            if self.local_rank == 0:
                time_start = time.time()
            for i, (feats, labs) in enumerate(self.loader):
                self.optimizer.zero_grad()
                feats = feats.reshape(-1, self.dataset.get_time_step(), \
                                      self.dataset.get_input_size())
                feats = feats.to(self.device)
                labs = labs.to(self.device)
                outputs = self.model(feats)
                loss = self.criterion(outputs, labs)
                loss.backward()
                self.optimizer.step()
            torch.cuda.synchronize()
            if self.local_rank == 0:
                print('e2e time. {}s'.format((time.time() - time_start) / len(self.loader)))
                
                t1 = time.time()
                dt[0] += t1 - t0

                cur_acc =self._validate2()
                if cur_acc>best_acc:
                    best_acc=cur_acc
                acc_str='Best_acc: {:.4f}'.format(best_acc)
                print(acc_str)
                fitlog.append(acc_str, with_time=True, change_line=True)
            # print(f'Pre Epoch. ({ self.local_rank} {t1-t0:.5f}s)')

            t2 = time.time()
            if lowest_loss > loss.item():
                lowest_loss = loss.item()
                lowest_loss_at = epoch + 1
            # print(f'Loss. ({time.time()-t2:.5f}s)')

            if i + 1 == nbatch and self.local_rank == 0:
                ptstr = 'Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}, Best: {:.4f}, Best at: {}'.format( \
                    epoch + 1, self.epochs, i + 1, nbatch, loss.item(), lowest_loss, lowest_loss_at)
                print(ptstr)
                fitlog.append(ptstr, with_time=True, change_line=True)
              
                
                # acc_list.append(acc)
                # val_loss_list.append(val_loss.cpu().numpy())
                # acc1=np.mean(acc_list)
                # val_loss1=np.mean(val_loss_list)
                # print("Epoch [{}/{}],acc:{},val_loss:{}".format(epoch + 1, self.epochs,str(acc1), str(val_loss1)),with_time=True, change_line=True)
                # fitlog.append("acc:{},val_loss:{}".format(str(acc1), str(val_loss1)),with_time=True, change_line=True)
                if loss.item() < self.early_stop_thresh:
                    early_stop_nreach += 1
                else:
                    early_stop_nreach = 0
                # prof.step()
            t3 = time.time()
            dt[1] += t3 - t2
            t4 = time.time()
            dt[2] += t4 - t3

            # ptstr, acc, val_loss=self.validate()
            # fitlog.append(ptstr+",acc:{},val_loss:{}".format(str(acc), str(val_loss)),with_time=True, change_line=True)

            # if early_stop_nreach >= self.early_stop_nreach_limit:
            #     break
        if self.local_rank == 0:
            fitlog.append(f'Done {self.epochs} Epoch. ({time.time() - ts:.5f}s)', with_time=True, change_line=True)
            fitlog.append("Train: {}s, Log: {}s Val: {}s".format(dt[0], dt[1], dt[2]), with_time=True, change_line=True)
            fitlog.close()
        print(f'Done {self.epochs} Epoch. ({time.time() - ts:.5f}s)')
        print("Train: {}s, Log: {}s Val: {}s".format(dt[0], dt[1], dt[2]))


   
    
    def _validate2(self):
        self.model.eval()
        self.dataset.set_mode('test')
        all_pred = []
        all_tar = []
        accs = []
        all_loss = []

        with torch.no_grad():
            for i, (ft, labs) in enumerate(self.loader):
                ft, labs = ft.to(self.device), labs.to(self.device)
                output = self.model(ft)
                loss = self.criterion(output, labs)
                preds = torch.argmax(output, dim=1).cpu().numpy().tolist()
                all_pred.extend(preds)
                all_tar.extend(labs.cpu().numpy().tolist())
                accs.append(accuracy_score(all_tar, all_pred))
                all_loss.append(loss.item())
                # if i % 100 == 0:
                #     print('validating @ batch {}'.format(i))

                # if g_dubug:
                #     break
        # ptstr = "Validation ACC: %.4f, loss: %.4f" % (np.mean(accs), np.mean(all_loss))
        
        # self.dataset.set_mode('train')
        return np.mean(accs)
        # return np.mean(accs), np.mean(all_loss),np.std(all_loss), all_pred, all_tar
        


def init_ddp(visiable_devices='0,1,2,3,4,5,6,7,8'):
    if torch.cuda.device_count() > 1:
        os.environ['HIP_VISIBLE_DEVICES'] = visiable_devices
        local_rank = int(os.environ["LOCAL_RANK"])
        print("local_rank:" + str(local_rank))
        #torch.distributed.init_process_group(backend='nccl', init_method='tcp://localhost:23456', rank=0, world_size=1)
        torch.distributed.init_process_group(backend="nccl")

        # local_rank = torch.distributed.get_rank()
        torch.cuda.set_device(local_rank)
        # device = torch.device("cuda", args.local_rank)
        return local_rank
    else:
        return None


if __name__ == "__main__":
    local_rank = init_ddp()
    driver = LSTMDriverHour(local_rank=local_rank)
    driver.load('dat_3day_pcase')
    driver.fit()