train.py - master - stupin/DreamNet

train.py
214 строк7 KB
Stupin
init
25 июн 2025, 19:54
25 июн 2025, 19:54b04fcff
from torch.cuda import is_available as cia
import torch.nn.functional as F
from App import Net
import numpy as np
import traceback
import datetime
import torch
import copy

class SparseDataDataset(torch.utils.data.Dataset):
    def __init__(self, data, targets):
        self.data = data
        self.targets = targets

    def __len__(self):
        return self.data.shape[0]
    
    def __getitem__(self, idx):
        cur_element = torch.from_numpy(self.data[idx]).float()
        cur_label = torch.from_numpy(np.asarray(self.targets[idx])).long()
        return cur_element, cur_label
    
def to_device(data, device):
    """Move tensor(s) to chosen device"""
    if isinstance(data, (list,tuple)):
        return [to_device(x, device) for x in data]
    return data.to(device, non_blocking=True)

class DeviceDataLoader():
    """Wrap a dataloader to move data to a device"""
    def __init__(self, dl, device):
        self.dl = dl
        self.device = device
        
    def __iter__(self):
        """Yield a batch of data after moving it to device"""
        for b in self.dl: 
            yield to_device(b, self.device)

    def __len__(self):
        """Number of batches"""
        return len(self.dl)
    
def get_raw_dataset() -> object:
    '''Return generator of raw data'''
    PATH = 'tweet_emotions.csv'
    with open(PATH, encoding='utf-8') as f:
        for line in f:
            line = line.split(',')
            yield [line[1], [line[2][_:_ + 3] for _ in range(len(line[2]) - 2)]]

def main() -> None:
    # Getting raw data
    raw = list(get_raw_dataset())
    
    # Forming and saving dictionaries
    lbl = list(set(data[0] for data in raw))
    dct = list(set(elem for data in raw for elem in data[1]))

    PATH = 'tweet_emotions_tokens.txt'
    with open(PATH, 'w', encoding='utf-8') as f:
        lbls = "\n".join(lbl)
        f.write(f'{"".join(dct)}\n{lbls}')

    # Founding friquency
    friq = {}

    for data in raw:
        for dt in data[1]:
            try:
                friq[dt] += 1
            except KeyError:
                friq[dt] = 1
    
    print(len(dct))
    for c, key in enumerate(friq.keys()):
        if c % 100 == 0:
            print(c)
        if friq[key] < MIN_COUNT:
            del dct[dct.index(key)]
    
    # Process data
    labels = np.array([], dtype=int)
    tokens = np.zeros((len(raw), len(dct)), dtype=np.float32)
    print(len(dct))
    for idx, data in enumerate(raw):
        if idx % 100 == 0:
            print(idx)
        labels = np.append(labels, lbl.index(data[0]))
        for elem in data[1]:
            try:
                tokens[idx][dct.index(elem)] = 1.0
            except ValueError:
                continue

    # Shuffle data
    p = np.random.permutation(labels.shape[0])
    tokens = tokens[p]
    labels = labels[p]

    # Split data to train and test datasets
    s = int(labels.shape[0] * SPLIT_COEF)

    trainset = SparseDataDataset(tokens[:s], labels[:s])
    trainloader = torch.utils.data.DataLoader(trainset,
                                              batch_size=BATCH_SIZE,
                                              shuffle=True,
                                              num_workers=2,
                                              pin_memory=cia())
    trainloader = DeviceDataLoader(trainloader, device)

    testset = SparseDataDataset(tokens[s:], labels[s:])
    testloader = torch.utils.data.DataLoader(testset,
                                             batch_size=BATCH_SIZE,
                                             shuffle=False,
                                             num_workers=2,
                                             pin_memory=torch.cuda.is_available())
    testloader = DeviceDataLoader(testloader, device)

    model = TextClassificationModel(len(dct), len(lbl))
    model = to_device(model, device)

    criterion = F.cross_entropy

    optimizer = torch.optim.Adam(model.parameters(),
                                 lr=1e-1,
                                 weight_decay=0)

    lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
                                                              patience=5,
                                                              factor=0.5,
                                                              verbose=True)
    
    best_test_loss = float('inf')
    best_epoch_i = 0
    best_model = copy.deepcopy(model)

    max_batches_per_epoch_train = 10000
    early_stopping_patience = 10

    for epoch_i in range(EPOCHS):
        try:
            epoch_start = datetime.datetime.now()
            print(f'Эпоха {epoch_i}')

            model.train()
            mean_train_loss = 0
            train_batches_n = 0
            for batch_i, (batch_x, batch_y) in enumerate(trainloader):
                if batch_i > max_batches_per_epoch_train:
                    break

                batch_x = to_device(batch_x, device)
                batch_y = to_device(batch_y, device)

                pred = model(batch_x)
                loss = criterion(pred, batch_y)

                model.zero_grad()
                loss.backward()
                optimizer.step()

                mean_train_loss += float(loss)
                train_batches_n += 1

            mean_train_loss /= train_batches_n
            print(f'Эпоха: {batch_i} итераций, {(datetime.datetime.now() - epoch_start).total_seconds()} сек')
            print('Среднее значение функции потерь на обучении', mean_train_loss)

            model.eval()
            mean_test_loss = 0
            test_batches_n = 0

            with torch.no_grad():
                for batch_i, (batch_x, batch_y) in enumerate(testloader):
                    if batch_i > max_batches_per_epoch_train:
                        break
                    
                    batch_x = to_device(batch_x, device)
                    batch_y = to_device(batch_y, device)

                    pred = model(batch_x)
                    loss = criterion(pred, batch_y)

                    mean_test_loss += float(loss)
                    test_batches_n += 1

            mean_test_loss /= test_batches_n
            print('Среднее значение функции потерь на тесте', mean_test_loss)

            if mean_test_loss < best_test_loss:
                best_epoch_i = epoch_i
                best_test_loss = mean_test_loss
                best_model = copy.deepcopy(model)
                print('Новая лучшая модель!')
            elif epoch_i - best_epoch_i > early_stopping_patience:
                print(f'Модель не улучшилась за последние {early_stopping_patience} эпох, прекращаем обучение')
                break

            lr_scheduler.step(mean_test_loss)

            print()
        except KeyboardInterrupt:
            print('Досрочно остановлено пользователем')
            break
        except Exception as e:
            print(f'Ошибка при обучении: {e}\n{traceback.format_exc()}')
            break
    
    torch.save(best_model.state_dict(), 'model.h5')
    torch.save(model.state_dict(), 'model_full.h5')

if __name__ == '__main__':
    main()