def train(args):
from train_utils import ModelTrainer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Prepare Data
train_dataset = Chime_Dataset('tr', args)
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
pin_memory=True,
collate_fn=lambda x: Chime_Collate(x),
num_workers=args.num_workers)
# Prepare model
if args.model_type == 'BLSTM':
model = BLSTMMaskEstimator()
model_save_dir = os.path.join(args.data_dir, 'BLSTM_model')
mkdir_p(model_save_dir)
elif args.model_type == 'FW':
model = SimpleFWMaskEstimator()
model_save_dir = os.path.join(args.data_dir, 'FW_model')
mkdir_p(model_save_dir)
else:
raise ValueError('Unknown model type. Possible are "BLSTM" and "FW"')
criterion = torch.nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters(),
lr=args.learning_rate,
weight_decay=1e-4)
trainer = ModelTrainer(model, train_loader, criterion, optimizer, args,
device)
trainer.train(args.num_epochs)
