Example #1
0
    def train(self, triplet=True):
        # generate dataset for PyTorch
        from data.dataset import TruncatedInputFromMFB, ToTensor, SpeakerDataset
        from torchvision import transforms
        from torch.utils.data import DataLoader
        import torch
        transform = transforms.Compose([TruncatedInputFromMFB(), ToTensor()])
        if Config.MODEL_TYPE == 'cnn3d':
            from data.dataset3d import SpeakerDataset3D
            initial_dataset = SpeakerDataset(transform=transform)
            train_dataset = SpeakerDataset3D(initial_dataset)
        else:
            train_dataset = SpeakerDataset(transform=transform)

        # instantiate a model
        if Config.MODEL_TYPE == 'rescnn':
            from models.rescnn import ResNet
            model_ = ResNet(layers=Config.RESCNN_LAYERS,
                            num_classes=Config.NUM_CLASSES)
        elif Config.MODEL_TYPE == 'gru':
            from models.gru import GRU
            model_ = GRU(layers=Config.GRU_LAYERS,
                         num_classes=Config.NUM_CLASSES)
        elif Config.MODEL_TYPE == 'cnn3d':
            from models.cnn3d import CNN3D
            model_ = CNN3D(num_classes=Config.NUM_CLASSES)

        from utils.train import Trainer
        model_ = model_.cuda()
        epoch = Config.SOFTMAX_TRAINING_EPOCH
        for i in range(epoch):
            optimizer = torch.optim.Adam(model_.parameters())
            train_loader = DataLoader(train_dataset,
                                      batch_size=Config.PRETRAIN_BATCH_SIZE,
                                      shuffle=True)
            Trainer.train(train_loader, model_, optimizer, i)

        if triplet:
            from copy import deepcopy
            model_tri = deepcopy(model_)
            model_tri = model_tri.cuda()
            epoch_ = Config.TRIPLET_TRAINING_EPOCH
            for i in range(epoch_):
                optimizer_ = torch.optim.SGD(model_tri.parameters(),
                                             lr=Config.TRIPLET_LR -
                                             i * Config.TRIPLET_LR_DECAY,
                                             momentum=Config.TRIPLET_MOMENTUM)
                train_loader = DataLoader(
                    train_dataset,
                    batch_size=Config.FINETUNE_BATCH_SIZE,
                    shuffle=True)
                Trainer.train_tri(train_loader,
                                  model_tri,
                                  optimizer_,
                                  i,
                                  semi_hard=True,
                                  triplet_margin=Config.TRIPLET_MARGIN)
def main():
    batch_size = args.batch_size
    augment = args.augment
    dataloaders = train_val_dataloaders(TRAIN_DIR, DEV_DIR, augment, batch_size)

    train_dataloader = dataloaders["train"]
    val_dataloader = dataloaders["val"]

    show = args.show_batch

    if show:
        # Let's have a look at the first batch
        print("Show batch from train dataloader: ")
        show_batch(train_dataloader, class_names)

        print("Show batch from val dataloader: ")
        show_batch(val_dataloader, class_names)

    # Model - pretrained ResNet18, trained on ImageNet
    model = models.resnet18(pretrained=True)

    # Disable grad for all conv layers
    for param in model.parameters():
        param.requires_grad = False
    print(
        "Output of ResNet18 before FC layer, that we add later: ", model.fc.in_features
    )
    # Add FC layer with 2 outputs: cleaned or dirty
    model.fc = torch.nn.Linear(model.fc.in_features, 5)

    # Put model on GPU
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    # Loss function - binary Cross-Entropy
    loss = torch.nn.CrossEntropyLoss()

    learning_rate = args.learning_rate
    # Optimization method - Adam
    optimizer = torch.optim.Adam(model.parameters(), amsgrad=True, lr=learning_rate)

    number_epochs = args.num_epochs
    # Training
    model_trainer = Trainer(
        model,
        train_dataloader,
        val_dataloader,
        loss,
        optimizer,
        device,
        num_epochs=number_epochs,
    )

    print("Begin training: ")
    model_trainer.train()

    # Draw losses and save the plot
    model_trainer.draw_losses("ResNet18", "Adam")

    # Save weights of the model
    model_trainer.save_weights("ResNet18", "Adam")

    # Save trained model
    model_trainer.save_model("ResNet18", "Adam")