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")