def train_autoencoder(device, args): # model definition model = FeatureExtractor() model.to(device) # data definition all_chunks = [] # concatenate all chunk files # note that it is independent of the # class of each chunk sinc we are creating # a generative dataset for label in filesystem.listdir_complete(filesystem.train_audio_chunks_dir): chunks = filesystem.listdir_complete(label) all_chunks = all_chunks + chunks train_chunks, eval_chunks = train_test_split(all_chunks, test_size=args.eval_size) # transforms and dataset trf = normalize train_dataset = GenerativeDataset(train_chunks, transforms=trf) eval_dataset = GenerativeDataset(eval_chunks, transforms=trf) train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=4, collate_fn=None,pin_memory=True) eval_dataloader = DataLoader(eval_dataset, batch_size=1, shuffle=True, num_workers=4, collate_fn=None,pin_memory=True) # main loop optimizer = optim.Adam(model.parameters(), lr=args.lr) loss_criterion = SoftDTW(use_cuda=True, gamma=0.1) train_count = 0 eval_count = 0 for epoch in range(args.n_epochs): print('Epoch:', epoch, '/', args.n_epochs) train_count = train_step(model, train_dataloader, optimizer, loss_criterion, args.verbose_epochs, device, train_count) eval_count = eval_step(model, eval_dataloader, loss_criterion, args.verbose_epochs, device, eval_count) torch.save(model.state_dict(), os.path.join(wandb.run.dir, 'model_checkpoint.pt'))
print( 'Test: D_Loss: {:.4f} F_Loss: {:.4f} C_Loss: {:.4f} Acc: {:.4f}' .format(D_loss, F_loss, C_loss, gesture_acc)) print('Time usage: {:.2f}s'.format(time.time() - epoch_start_time)) # scheduler.step(F_loss) scheduler_F.step(F_loss) scheduler_D.step(D_loss) scheduler_C.step(C_loss) if F_loss + precision < best_loss: print('New best validation F_loss: {:.4f}'.format(F_loss)) best_loss = F_loss best_weights_F = copy.deepcopy(feature_extractor.state_dict()) best_weights_C = copy.deepcopy(label_predictor.state_dict()) best_weights_D = copy.deepcopy(domain_classifier.state_dict()) patience = patience_increase + epoch print('So Far Patience: ', patience) print() # save model torch.save(best_weights_F, r'saved_model\feature_extractor.pkl') torch.save(best_weights_C, r'saved_model\label_predictor.pkl') torch.save(best_weights_D, r'saved_model\domain_classifier.pkl') print('Model Saved.') # torch.save(feature_extractor, r'saved_model\m_feature_extractor.pkl') # torch.save(label_predictor, r'saved_model\m_label_predictor.pkl')
torch.argmax(class_logits, dim=1) == source_label).item() total_num += source_data.shape[0] print(i, end='\r') return running_D_loss / (i + 1), running_F_loss / ( i + 1), total_hit / total_num # 訓練200 epochs for epoch in range(1000): train_D_loss, train_F_loss, train_acc = train_epoch( source_dataloader, target_dataloader, lamb=((2 / (1 + math.exp(-10 * epoch / 999))) - 1)) if (epoch + 1) % 100 == 0: torch.save(feature_extractor.state_dict(), f'_semi_lamb_extractor_model_{epoch+1:03d}.bin') torch.save(label_predictor.state_dict(), f'_semi_lamb_predictor_model_{epoch+1:03d}.bin') print( 'epoch {:>3d}: train D loss: {:6.4f}, train F loss: {:6.4f}, acc {:6.4f}' .format(epoch, train_D_loss, train_F_loss, train_acc)) # test result = [] label_predictor.eval() feature_extractor.eval() for i, (test_data, _) in enumerate(test_dataloader): test_data = test_data.cuda() class_logits = label_predictor(feature_extractor(test_data))
def main(num_epochs=10, embedding_dim=256, data_dir="data/"): """ Function to train the model. Args: num_epochs: int Number of full dataset iterations to train the model. embedding_dim: int Output of the CNN model and input of the LSTM embedding size. data_dir: str Path to the folder of the data. """ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(f"WORKING WITH: {device}") # Define the paths for train and validation train_json_path = data_dir + "annotations/captions_train2014.json" train_root_dir = data_dir + "train2014" valid_json_path = data_dir + "annotations/captions_val2014.json" valid_root_dir = data_dir + "val2014" transform = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) train_dataset = CocoDataset(json_path=train_json_path, root_dir=train_root_dir, transform=transform) train_coco_dataset = get_data_loader(train_dataset, batch_size=128) valid_dataset = CocoDataset(json_path=valid_json_path, root_dir=valid_root_dir, transform=transform) valid_coco_dataset = get_data_loader(valid_dataset, batch_size=1) encoder = FeatureExtractor(embedding_dim).to(device) decoder = CaptionGenerator(embedding_dim, 512, len(train_dataset.vocabulary), 1).to(device) criterion = nn.CrossEntropyLoss() # params = list(decoder.parameters()) + list(encoder.linear.parameters()) + list(encoder.bn.parameters()) params = list(decoder.parameters()) + list( encoder.linear.parameters()) + list(encoder.bn.parameters()) optimizer = optim.Adam(params, lr=0.01) print(f"TRAIN DATASET: {len(train_coco_dataset)}") print(f"VALID DATASET: {len(valid_coco_dataset)}") total_step = len(train_coco_dataset) for epoch in range(num_epochs): encoder.train() decoder.train() train_loss = 0.0 valid_loss = 0.0 for i, (images, captions, descriptions) in enumerate(train_coco_dataset): # targets = pack_padded_sequence(caption, 0, batch_first=True)[0] images = images.to(device) captions = captions.to(device) # targets = pack_padded_sequence(captions, lengths, batch_first=True)[0] features = encoder(images) outputs = decoder(features, captions) loss = criterion(outputs.view(-1, len(train_dataset.vocabulary)), captions.view(-1)) # bleu = calculate_bleu(decoder, features, descriptions, coco_dataset) # print(bleu) encoder.zero_grad() decoder.zero_grad() loss.backward() optimizer.step() # Print log info train_loss += loss.item() ''' if i % 10 == 0: print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}, Perplexity: {:5.4f}' .format(epoch, num_epochs, i, total_step, loss.item(), np.exp(loss.item()))) ''' # Save the model checkpoints if (i + 1) % 1000 == 0: torch.save( decoder.state_dict(), os.path.join("models", 'decoder-{}-{}.ckpt'.format(epoch + 1, i + 1))) torch.save( encoder.state_dict(), os.path.join("models", 'encoder-{}-{}.ckpt'.format(epoch + 1, i + 1))) encoder.eval() decoder.eval() bleu = 0.0 for i, (images, captions, descriptions) in enumerate(valid_coco_dataset): if (i > 80000): break images = images.to(device) captions = captions.to(device) features = encoder(images) outputs = decoder(features, captions) loss = criterion(outputs.view(-1, len(train_dataset.vocabulary)), captions.view(-1)) valid_loss += loss.item() bleu += calculate_bleu(decoder, features, descriptions, train_coco_dataset) # print(f"BLEU: {bleu / 10000}") print( "Epoch: {}, Train Loss: {:.4f}, Valid Loss: {:.4f}, BLEU: {:.4f}". format(epoch, train_loss / len(train_coco_dataset), valid_loss / 80000, bleu / 80000))