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'))
def train_classifier(device, args): encoder = FeatureExtractor() encoder.load_state_dict(torch.load(args.encoder_path)) encoder.eval() classifier = Classifier(encoder) classifier.to(device) all_chunks = [] all_labels = [] for label in filesystem.listdir_complete(filesystem.train_audio_chunks_dir): chunks = filesystem.listdir_complete(label) all_chunks = all_chunks + chunks all_labels = all_labels + [label.split('/')[-1]] * len(chunks) train_chunks, eval_chunks, train_labels, eval_labels = train_test_split(all_chunks, all_labels, test_size=args.eval_size) # transforms and dataset trf = normalize # dataset generation labels_encoder = LabelsEncoder(pd.read_csv(filesystem.labels_encoding_file)) train_dataset = DiscriminativeDataset(train_chunks, train_labels, labels_encoder, transforms=trf) eval_dataset = DiscriminativeDataset(eval_chunks, eval_labels, labels_encoder, 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) optimizer = optim.Adam(classifier.parameters(), lr=args.lr) loss_criterion = nn.CrossEntropyLoss() train_count = 0 eval_count = 0 for epoch in range(args.n_epochs): print('Epoch:', epoch, '/', args.n_epochs) train_count = train_step_classification(classifier, train_dataloader, optimizer, loss_criterion, args.verbose_epochs, device, train_count) torch.save(classifier.state_dict(), os.path.join(wandb.run.dir, 'model_checkpoint.pt')) eval_count = eval_step_classification(classifier, eval_dataloader, loss_criterion, args.verbose_epochs, device, eval_count)
def main(): # define parameters num_class = 6 batch_size = 1 time_step = 32 cnn_feat_size = 256 # AlexNet gaze_size = 3 gaze_lstm_hidden_size = 64 gaze_lstm_projected_size = 128 # dataset_path = '../data/gaze_dataset' dataset_path = '../../gaze-net/gaze_dataset' img_size = (224, 224) time_skip = 2 # define model arch = 'alexnet' extractor_model = FeatureExtractor(arch=arch) extractor_model.features = torch.nn.DataParallel(extractor_model.features) extractor_model.cuda() # uncomment this line if using cpu extractor_model.eval() model = SpatialAttentionModel(num_class, cnn_feat_size, gaze_size, gaze_lstm_hidden_size, gaze_lstm_projected_size) model.cuda() # load model from checkpoint model = load_checkpoint(model) trainGenerator = gaze_gen.GazeDataGenerator(validation_split=0.2) train_data = trainGenerator.flow_from_directory(dataset_path, subset='training', crop=False, batch_size=batch_size, target_size= img_size, class_mode='sequence_pytorch', time_skip=time_skip) # small dataset, error using validation split val_data = trainGenerator.flow_from_directory(dataset_path, subset='validation', crop=False, batch_size=batch_size, target_size= img_size, class_mode='sequence_pytorch', time_skip=time_skip) # start predict for i in range(10): print("start a new interaction") # img_seq: (ts,224,224,3), gaze_seq: (ts, 3), ouput: (ts, 6) # [img_seq, gaze_seq], target = next(val_data) [img_seq, gaze_seq], target = next(train_data) restart = True predict(img_seq, gaze_seq, extractor_model, model, restart=restart) print(target) for j in range(img_seq.shape[0]): # predict(img_seq[j], gaze_seq[j], None, model, restart=restart) # print(target[j]) # restart = False img = img_seq[j,:,:,:] gazes = gaze_seq cv2.circle(img, (int(gazes[j,1]), int(gazes[j,2])), 10, (255,0,0),-1) cv2.imshow('ImageWindow', img) cv2.waitKey(33)
def main(): args = get_args() device = Config.device print("PyTorch running with device {0}".format(device)) print("Creating transformation network...") network = TransformationNetwork().to(device) if args.test: assert args.model_path, '--model_path must be provided for testing' print('Testing...') network.eval() print("Loading transformation network...") load_network(network, args.model_path) test_image_loader = load_image_dataloader(args.test_image_path, batch_size=Config.batch_size * 2, shuffle=False) if not os.path.isdir(args.image_save_path): os.makedirs(args.image_save_path) print("Generating images...") generate_and_save_images(network, test_image_loader, args.image_save_path) else: # Training! print("Training...") print("Loading feature extractor") feature_extractor = FeatureExtractor() # load dataloader print("Loading data") content_image_dataloader = load_image_dataloader(args.content_image_path, batch_size=args.batch_size) style_image = load_style_image(args.style_image_path) print("Loading Trainer") trainer = Trainer(network, feature_extractor, content_image_dataloader, style_image, lr=args.lr) if args.model_path: trainer.load_checkpoint(args.model_path) print("Start training") loss_hist = trainer.train(num_epochs=args.num_epochs) if not os.path.isdir(args.model_save_path): os.makedirs(args.model_save_path) trainer.save_checkpoint(os.path.join(args.model_save_path, 'epoch-{0}.ckpt'.format(args.num_epochs))) plt.plot(loss_hist, label='Loss') plt.legend() plt.savefig('StyleTransfer_train_history.jpg') plt.show()
def main(args): # Image preprocessing transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) image_dir = "data/" json_path = image_dir + "annotations/captions_train2014.json" root_dir = image_dir + "train2014" dataset = CocoDataset(json_path=json_path, root_dir=root_dir, transform=transform) data_loader = get_data_loader(dataset, batch_size=32) # Build models encoder = FeatureExtractor(args.embed_size).eval( ) # eval mode (batchnorm uses moving mean/variance) decoder = CaptionGenerator(args.embed_size, args.hidden_size, len(dataset.vocabulary), args.num_layers) encoder = encoder.to(device) decoder = decoder.to(device) # Load the trained model parameters encoder.load_state_dict(torch.load(args.encoder_path)) decoder.load_state_dict(torch.load(args.decoder_path)) # Prepare an image image = load_image(args.image, transform) image_tensor = image.to(device) # Generate an caption from the image feature = encoder(image_tensor) sampled_ids = decoder.sample(feature) sampled_ids = sampled_ids[0].cpu().numpy( ) # (1, max_seq_length) -> (max_seq_length) # Convert word_ids to words sampled_caption = [] for word_id in sampled_ids: word = data_loader.dataset.id_to_word[word_id] sampled_caption.append(word) if word == '<end>': break sentence = ' '.join(sampled_caption) # Print out the image and the generated caption print(sentence) image = Image.open(args.image) plt.imshow(np.asarray(image))
def __init__(self, src_domain, tgt_domain): self.num_epoch = 10 self.gamma = 1.0 print('construct dataset and dataloader...') train_dataset = TrainDataset(src_domain, tgt_domain) self.NEG_NUM = train_dataset.NEG_NUM self.input_dim = train_dataset.sample_dim self.train_loader = DataLoader(train_dataset, batch_size=32) print('Done!') self.feature_extractor = FeatureExtractor(self.input_dim) self.optimizer = optim.SGD(self.feature_extractor.parameters(), lr=0.1, momentum=0.9)
def main(): parser = get_parser() args = parser.parse_args() model_path = args.model input_path = args.input sound_path = args.output model = FeatureExtractor() model.load_state_dict(torch.load(model_path)) device = torch.device('cuda') cpu_device = torch.device('cpu') model.to(device) #data = normalize(torchaudio.load(input_path)[0][0].reshape(1, -1)) data = torch.from_numpy(normalize(torch.randn(1, 132480))).float().to(device) data = data.reshape(1, 1, -1) model.eval() sound = model(data) print(functional.mse_loss(sound, data).item()) sound = sound.to(cpu_device) torchaudio.save(sound_path, sound.reshape(-1), 44100)
def model_env(args): """building model environment avoiding to instantiate model. Args: args : model arguments which is control by demo_utils.argument_setting Returns: model (torch.nn): build model in cuda device criterion(torch.nn): build criterion. Default to mse loss extractor(torch.nn): build vgg content loss in cuda device """ if args.doc: args = config_loader(args.doc, args) # set cuda device torch.cuda.set_device(args.gpu_id) # model version control version = args.load if type(args.load) is int else 0 # model path and parameter model_path = os.path.join( args.log_path, args.model_name, f'version_{version}',f'{args.model_name}_{args.scale}x.pt') checkpoint = torch.load(model_path, map_location=f'cuda:{args.gpu_id}') # loading model model = model_builder(args.model_name, args.scale, **args.model_args).cuda() model.load_state_dict(checkpoint['state_dict']) # build criterion criterion = criterion_builder(args.criterion) # loading feature extractor extractor = FeatureExtractor().cuda() if args.content_loss else None return model, criterion, extractor
def train(model, train_loader, valid_loader, optimizer, criterion, args): # content_loss best_err = None feature_extractor = FeatureExtractor().cuda() feature_extractor.eval() writer, log_path = writer_builder(args.log_path, args.model_name, load=args.load) # init data checkpoint = { 'epoch': 1, # start from 1 'train_iter': 0, # train iteration 'valid_iter': 0, # valid iteration } model_path = os.path.join(log_path, f'{args.model_name}_{args.scale}x.pt') # config model_config(train_args, save=log_path) # save model configuration before training # load model from exist .pt file if args.load and os.path.isfile(model_path): r""" load a pickle file from exist parameter state_dict: model's state dict epoch: parameters were updated in which epoch """ checkpoint = torch.load(model_path, map_location=f'cuda:{args.gpu_id}') checkpoint['epoch'] += 1 # start from next epoch checkpoint['train_iter'] += 1 checkpoint['valid_iter'] += 1 model.load_state_dict(checkpoint['state_dict']) # initialize the early_stopping object if args.early_stop: early_stopping = EarlyStopping(patience=args.patience, threshold=args.threshold, verbose=args.verbose, path=model_path) if args.scheduler: scheduler = schedule_builder(optimizer, args.scheduler, args.step, args.factor) # progress bar postfix value pbar_postfix = { 'MSE loss': 0.0, 'Content loss': 0.0, 'lr': args.lr, } for epoch in range(checkpoint['epoch'], args.epochs + 1): model.train() err = 0.0 valid_err = 0.0 train_bar = tqdm(train_loader, desc=f'Train epoch: {epoch}/{args.epochs}') for data in train_bar: # load data from data loader inputs, target, _ = data inputs, target = inputs.cuda(), target.cuda() # predicted fixed 6 axis data pred = model(inputs) # MSE loss mse_loss = args.alpha * criterion(pred - inputs, target - inputs) # content loss gen_features = feature_extractor(pred) real_features = feature_extractor(target) content_loss = args.beta * criterion(gen_features, real_features) # for compatible but bad for memory usage loss = mse_loss + content_loss # update progress bar pbar_postfix['MSE loss'] = mse_loss.item() pbar_postfix['Content loss'] = content_loss.item() # show current lr if args.scheduler: pbar_postfix['lr'] = optimizer.param_groups[0]['lr'] train_bar.set_postfix(pbar_postfix) err += loss.sum().item() * inputs.size(0) optimizer.zero_grad() loss.backward() optimizer.step() # update writer writer.add_scalar('Iteration/train loss', loss.sum().item(), checkpoint['train_iter']) checkpoint['train_iter'] += 1 # cross validation valid_bar = tqdm(valid_loader, desc=f'Valid epoch:{epoch}/{args.epochs}', leave=False) model.eval() input_epoch = pred_epoch = target_epoch = torch.empty(0, 0) with torch.no_grad(): for data in valid_bar: # for data in valid_loader: inputs, target, _ = data inputs, target = inputs.cuda(), target.cuda() pred = model(inputs) # MSE loss mse_loss = criterion(pred - inputs, target - inputs) # content loss gen_features = feature_extractor(pred) real_features = feature_extractor(target) content_loss = criterion(gen_features, real_features) # for compatible loss = mse_loss + content_loss # update progress bar pbar_postfix['MSE loss'] = mse_loss.item() pbar_postfix['Content loss'] = content_loss.item() # show current lr if args.scheduler: pbar_postfix['lr'] = optimizer.param_groups[0]['lr'] valid_bar.set_postfix(pbar_postfix) valid_err += loss.sum().item() * inputs.size(0) # update writer writer.add_scalar('Iteration/valid loss', loss.sum().item(), checkpoint['valid_iter']) checkpoint['valid_iter'] += 1 # out2csv every check interval epochs (default: 5) if epoch % args.check_interval == 0: input_epoch = inputs pred_epoch = pred target_epoch = target # out2csv every check interval epochs (default: 5) if epoch % args.check_interval == 0: # tensor to csv file out2csv(input_epoch, f'{epoch}', 'input', args.out_num, args.save_path, args.stroke_length) out2csv(pred_epoch, f'{epoch}', 'output', args.out_num, args.save_path, args.stroke_length) out2csv(target_epoch, f'{epoch}', 'target', args.out_num, args.save_path, args.stroke_length) # compute loss err /= len(train_loader.dataset) valid_err /= len(valid_loader.dataset) print(f'\ntrain loss: {err:.4f}, valid loss: {valid_err:.4f}') # update scheduler if args.scheduler: scheduler.step() # update loggers writer.add_scalars( 'Epoch', { 'train loss': err, 'valid loss': valid_err }, epoch, ) # early_stopping needs the validation loss to check if it has decresed, # and if it has, it will make a checkpoint of the current model if args.early_stop: early_stopping(valid_err, model, epoch) if early_stopping.early_stop: print("Early stopping") break # if early stop is false, store model when the err is lowest elif epoch == checkpoint['epoch'] or err < best_err: best_err = err # save err in first epoch # save current epoch and model parameters torch.save( { 'state_dict': model.state_dict(), 'epoch': epoch, 'train_iter': checkpoint['train_iter'], 'valid_iter': checkpoint['valid_iter'], }, model_path) writer.close()
def get_network(): return FeatureExtractor(), DifferenceClassifier()
def main(): # define parameters TRAIN = True num_class = 6 batch_size = 1 # time_step = 32 epochs = 50 cnn_feat_size = 256 # AlexNet gaze_size = 3 gaze_lstm_hidden_size = 64 gaze_lstm_projected_size = 128 learning_rate = 0.0001 momentum = 0.9 weight_decay = 1e-4 eval_freq = 1 # epoch print_freq = 1 # iteration # dataset_path = '../data/gaze_dataset' dataset_path = '../../gaze-net/gaze_dataset' img_size = (224, 224) log_path = '../log' logger = Logger(log_path, 'spatial') # define model arch = 'alexnet' extractor_model = FeatureExtractor(arch=arch) extractor_model.features = torch.nn.DataParallel(extractor_model.features) extractor_model.cuda() # uncomment this line if using cpu extractor_model.eval() model = SpatialAttentionModel(num_class, cnn_feat_size, gaze_size, gaze_lstm_hidden_size, gaze_lstm_projected_size) model.cuda() # define loss and optimizer # criterion = nn.CrossEntropyLoss() criterion = nn.CrossEntropyLoss().cuda() optimizer = torch.optim.SGD(model.parameters(), learning_rate, momentum = momentum, weight_decay=weight_decay) # define generator trainGenerator = gaze_gen.GazeDataGenerator(validation_split=0.2) train_data = trainGenerator.flow_from_directory(dataset_path, subset='training', crop=False, batch_size=batch_size, target_size= img_size, class_mode='sequence_pytorch') # small dataset, error using validation split val_data = trainGenerator.flow_from_directory(dataset_path, subset='validation', crop=False, batch_size=batch_size, target_size= img_size, class_mode='sequence_pytorch') # val_data = train_data def test(train_data): [img_seq, gaze_seq], target = next(train_data) img = img_seq[100,:,:,:] img_gamma = adjust_contrast(img) imsave('contrast.jpg', img_gamma) imsave('original.jpg', img) # test(train_data) print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~") # img_seq: (ts,224,224,3), gaze_seq: (ts, 3), ouput: (ts, 6) # [img_seq, gaze_seq], output = next(train_data) # print("gaze data shape") # print(img_seq.shape) # print(gaze_seq.shape) # print(output.shape) # start Training para = {'bs': batch_size, 'img_size': img_size, 'num_class': num_class, 'print_freq': print_freq} if TRAIN: print("get into training mode") best_acc = 0 for epoch in range(epochs): adjust_learning_rate(optimizer, epoch, learning_rate) print('Epoch: {}'.format(epoch)) # train for one epoch train(train_data, extractor_model, model, criterion, optimizer, epoch, logger, para) # evaluate on validation set if epoch % eval_freq == 0 or epoch == epochs - 1: acc = validate(val_data, extractor_model, model, criterion, epoch, logger, para, False) is_best = acc > best_acc best_acc = max(acc, best_acc) save_checkpoint({ 'epoch': epoch + 1, 'arch': arch, 'state_dict': model.state_dict(), 'best_acc': best_acc, 'optimizer': optimizer.state_dict(), }, is_best) else: model = load_checkpoint(model) print("get into testing and visualization mode") print("visualization for training data") vis_data_path = '../vis/train/' if not os.path.exists(vis_data_path): os.makedirs(vis_data_path) acc = validate(train_data, extractor_model, model, criterion, -1, \ logger, para, False, vis_data_path) print("visualization for validation data") vis_data_path = '../vis/val/' if not os.path.exists(vis_data_path): os.makedirs(vis_data_path) acc = validate(val_data, extractor_model, model, criterion, -1, \ logger, para, True, vis_data_path)
def main(): # define parameters TRAIN = True time_skip = 2 num_class = 6 batch_size = 1 epochs = 50 cnn_feat_size = 256 # AlexNet gaze_size = 3 gaze_lstm_hidden_size = 64 gaze_lstm_projected_size = 128 temporal_projected_size = 128 queue_size = 32 learning_rate = 0.0001 momentum = 0.9 weight_decay = 1e-4 eval_freq = 1 # epoch print_freq = 1 # iteration dataset_path = '../../gaze-net/gaze_dataset' # dataset_path = '../../gaze-net/gaze_dataset' img_size = (224, 224) extractor = True # fine-tune the last two layers of feat_extractor or not log_path = '../log' logger = Logger(log_path, 'multiple') # define model if extractor == False: arch = 'alexnet' extractor_model = FeatureExtractor(arch=arch) extractor_model.features = torch.nn.DataParallel( extractor_model.features) # extractor_model.cuda() # uncomment this line if using cpu extractor_model.eval() else: extractor_model = None model = MultipleAttentionModel(num_class, cnn_feat_size, gaze_size, gaze_lstm_hidden_size, gaze_lstm_projected_size, temporal_projected_size, queue_size, extractor=extractor) model.cuda() # define loss and optimizer criterion = nn.CrossEntropyLoss().cuda() param_list = [] for i, param in enumerate(model.parameters()): if param.requires_grad == True: print(param.size()) param_list.append(param) optimizer = torch.optim.SGD(param_list, learning_rate, momentum=momentum, weight_decay=weight_decay) # define generator trainGenerator = gaze_gen.GazeDataGenerator(validation_split=0.2) train_data = trainGenerator.flow_from_directory( dataset_path, subset='training', crop=False, batch_size=batch_size, target_size=img_size, class_mode='sequence_pytorch', time_skip=time_skip) # small dataset, error using validation split val_data = trainGenerator.flow_from_directory( dataset_path, subset='validation', crop=False, batch_size=batch_size, target_size=img_size, class_mode='sequence_pytorch', time_skip=time_skip) # val_data = train_data # start Training para = { 'bs': batch_size, 'img_size': img_size, 'num_class': num_class, 'print_freq': print_freq } if TRAIN: print("get into training mode") best_acc = 0 # acc = validate(val_data, extractor_model, model, criterion, 0, logger, para, False) for epoch in range(epochs): adjust_learning_rate(optimizer, epoch, learning_rate) print('Epoch: {}'.format(epoch)) # train for one epoch train(train_data, extractor_model, model, criterion, optimizer, epoch, logger, para) # evaluate on validation set if epoch % eval_freq == 0 or epoch == epochs - 1: acc = validate(val_data, extractor_model, model, criterion, epoch, logger, para, False) is_best = acc > best_acc best_acc = max(acc, best_acc) save_checkpoint( { 'epoch': epoch + 1, 'arch': arch, 'state_dict': model.state_dict(), 'best_acc': best_acc, 'optimizer': optimizer.state_dict(), }, is_best) else: model = load_checkpoint(model) print("get into testing and visualization mode") print("visualization for training data") vis_data_path = '../vis/train/' if not os.path.exists(vis_data_path): os.makedirs(vis_data_path) acc = validate(train_data, extractor_model, model, criterion, -1, \ logger, para, True, vis_data_path) print("visualization for validation data") vis_data_path = '../vis/val/' if not os.path.exists(vis_data_path): os.makedirs(vis_data_path) acc = validate(val_data, extractor_model, model, criterion, -1, \ logger, para, True, vis_data_path)
def test(model, test_loader, criterion, args): number = args.test_path.split('test_')[1] save_path = './output/char0' + number print(f'Saving data in {save_path}') # set model path if args.load is not False: _, log_path = writer_builder(args.log_path, args.model_name, load=args.load) model_path = os.path.join(log_path, f'{args.model_name}_{args.scale}x.pt') # load model parameters checkpoint = torch.load(model_path, map_location=f'cuda:{args.gpu_id}') # try-except to compatible try: model.load_state_dict(checkpoint['state_dict']) except: print('Warning: load older version') model.feature = nn.Sequential(*model.feature, *model.bottle) model.bottle = nn.Sequential() model.load_state_dict(checkpoint['state_dict'], strict=False) model.eval() # declare content loss feature_extractor = FeatureExtractor().cuda() feature_extractor.eval() err = 0.0 # out2csv i = 0 # count the number of loops j = 0 # count the number of data for data in tqdm(test_loader, desc=f'scale: {args.scale}'): inputs, target, _ = data inputs, target = inputs.cuda(), target.cuda() # normalize inputs and target # inputs = input_scaler.fit(inputs) # target = target_scaler.fit(target) pred = model(inputs) # denormalize # pred = input_scaler.inverse_transform(pred) # out2csv while j - (i * args.batch_size) < pred.size(0): out2csv(inputs, f'test_{int(j/args.test_num)+1}', 'input', j - (i * args.batch_size), save_path, args.stroke_length, spec_flag=True) out2csv(pred, f'test_{int(j/args.test_num)+1}', 'output', j - (i * args.batch_size), save_path, args.stroke_length, spec_flag=True) out2csv(target, f'test_{int(j/args.test_num)+1}', 'target', j - (i * args.batch_size), save_path, args.stroke_length, spec_flag=True) j += args.test_num i += 1 # MSE loss mse_loss = args.alpha * criterion(pred - inputs, target - inputs) # content loss gen_feature = feature_extractor(pred) real_feature = feature_extractor(target) content_loss = args.beta * criterion(gen_feature, real_feature) # for compatible loss = content_loss + mse_loss err += loss.sum().item() * inputs.size(0) err /= len(test_loader.dataset) print(f'test error:{err:.4f}')
def get_network(): return FeatureExtractor()
def train(model, train_loader, valid_loader, optimizer, criterion, args): # declare content loss best_err = None feature_extractor = FeatureExtractor().cuda() feature_extractor.eval() # load data model_path = f'fsrcnn_{args.scale}x.pt' checkpoint = {'epoch': 1} # start from 1 # load model from exist .pt file if args.load is True and os.path.isfile(model_path): r""" load a pickle file from exist parameter state_dict: model's state dict epoch: parameters were updated in which epoch """ checkpoint = torch.load(model_path, map_location=f'cuda:{args.gpu_id}') checkpoint['epoch'] += 1 # start from next epoch model.load_state_dict(checkpoint['state_dict']) # store the training time writer = writer_builder(args.log_path,args.model_name) for epoch in range(checkpoint['epoch'], args.epochs+1): model.train() err = 0.0 valid_err = 0.0 store_data_cnt = 0 # to create new dataset for data in tqdm(train_loader, desc=f'train epoch: {epoch}/{args.epochs}'): # read data from data loader inputs, target, stroke_num = data inputs, target = inputs.cuda(), target.cuda() # predicted fixed 6 axis data pred = model(inputs) # inverse transform pred = inverse_scaler_transform(pred, inputs) # MSE loss mse_loss = criterion(pred, target) # content loss gen_features = feature_extractor(pred) real_features = feature_extractor(target) content_loss = criterion(gen_features, real_features) # for compatible but bad for memory usage loss = mse_loss + content_loss err += loss.sum().item() * inputs.size(0) # out2csv every check interval epochs (default: 5) if epoch % args.check_interval == 0: out2csv(inputs, f'{epoch}_input', args.stroke_length) out2csv(pred, f'{epoch}_output', args.stroke_length) out2csv(target, f'{epoch}_target', args.stroke_length) if epoch == args.epochs: if not os.path.exists('final_output'): os.mkdir('final_output') save_final_predict_and_new_dataset(pred, stroke_num, f'final_output/', args, store_data_cnt) store_data_cnt+=args.batch_size optimizer.zero_grad() loss.backward() optimizer.step() # cross validation model.eval() with torch.no_grad(): for data in tqdm(valid_loader, desc=f'valid epoch: {epoch}/{args.epochs}'): inputs, target = data inputs, target = inputs.cuda(), target.cuda() pred = model(inputs) # inverse transform pred = inverse_scaler_transform(pred, inputs) # MSE loss mse_loss = criterion(pred, target) # content loss gen_features = feature_extractor(pred) real_features = feature_extractor(target) content_loss = criterion(gen_features, real_features) # for compatible loss = mse_loss + content_loss valid_err += loss.sum().item() * inputs.size(0) if epoch == args.epochs: save_final_predict_and_new_dataset(pred, stroke_num, f'final_output/', args, store_data_cnt) store_data_cnt+=args.batch_size # compute loss err /= len(train_loader.dataset) valid_err /= len(valid_loader.dataset) print(f'train loss: {err:.4f}, valid loss: {valid_err:.4f}') # update every epoch # save model as pickle file if best_err is None or err < best_err: best_err = err # save current epoch and model parameters torch.save( { 'state_dict': model.state_dict(), 'epoch': epoch, } , model_path) # update loggers writer.add_scalars('Loss/', {'train loss': err, 'valid loss': valid_err}, epoch) writer.close()
# 'polar': {'max_len': vir_max_len, 'transform': trans_polar} } device = 'cuda' if torch.cuda.is_available() else 'cpu' dataset = Modalities( 'Exp0', split_cycle=split_cycle, start_date=start_date, end_date=end_date, **modalities ) train_amount = int(train_ratio * dataset.num_plants) test_amount = dataset.num_plants - train_amount train_set, test_set = ModalitiesSubset.random_split(dataset, [train_amount, test_amount]) train_loader = data.DataLoader(train_set, batch_size=batch_size, num_workers=2, shuffle=True) feat_ext = FeatureExtractor(*modalities).to(device) label_cls = nn.Sequential(nn.ReLU(), nn.Linear(512, len(classes)).to(device)) plant_cls = nn.Sequential(nn.ReLU(), nn.Linear(512, train_set.num_plants).to(device)) criterion = nn.CrossEntropyLoss().cuda() label_opt = optim.Adam(label_cls.parameters(), lr=label_lr) plant_opt = optim.Adam(plant_cls.parameters(), lr=plant_lr) ext_opt = optim.Adam(feat_ext.parameters(), lr=extractor_lr) best_loss = float('inf') def test_model(save_checkpoints=True): test_loader = data.DataLoader(test_set, batch_size=batch_size, num_workers=2, shuffle=True)
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))
learning_rate = config['learningRate'] style_weight = config['styleWeight'] content_weight = config['contentWeight'] tv_weight = config['tvWeight'] dataset = tf.data.Dataset.from_generator( image_loader, tf.float32, args=[config['datasetPath'], (256, 256, 3)]) optimizer = tf.optimizers.Adam(learning_rate=learning_rate) data_size_per_epoch = 4e4 step = data_size_per_epoch * epoch // batch_size auto_save_step = step // 10 transfer_model = TransferModel() feature_extractor = FeatureExtractor() style_targets = feature_extractor(tf.constant(style_image * 255.0))['style'] style_loss_metric = keras.metrics.Mean(name='style_loss') content_loss_metric = keras.metrics.Mean(name='content_loss') tv_loss_metric = keras.metrics.Mean(name='tv_loss') summary_writer = tf.summary.create_file_writer( os.path.join(model_path, 'logs')) bar = get_progress_bar() @tf.function() def train_step(X): with tf.GradientTape() as tape:
print ('loading weights of the generator ...') generator.load_state_dict(torch.load(args.gen_w)) print (generator) discriminator = Discriminator() if args.dis_w != '': print ('loading weights of the discriminator ...') discriminator.load_state_dict(torch.load(args.dis_w)) print (discriminator) # Put both G & D to device: generator = generator.to(device) discriminator = discriminator.to(device) # Loss criterion: encoding = FeatureExtractor(torchvision.models.vgg19(pretrained=True)).to(device) # VGG19 encoding; content_criterion = nn.MSELoss().to(device) adversarial_criterion = nn.BCELoss().to(device) ones = torch.nn.Parameter(torch.ones(args.batch_size), requires_grad=False).to(device) zeros = torch.nn.Parameter(torch.zeros(args.batch_size), requires_grad=False).to(device) # tensor placeholders ''' high_res = torch.FloatTensor(args.batch_size, 3, 96, 96) high_res_fake = torch.FloatTensor(args.batch_size, 3, 96, 96) output_fake = torch.FloatTensor(args.batch_size) output_real = torch.FloatTensor(args.batch_size) low_res = torch.FloatTensor(args.batch_size, 3, 24, 24).to(device) high_res = Variable(high_res) high_res_fake = Variable(high_res_fake) output_fake = Variable(output_fake)
def main(args: argparse.Namespace): checkpoint_name = get_checkpoint_name(args.experiment, args.excluded_modalities) # training hyper-parameters epochs = args.epochs batch_size = args.batch_size label_lr = args.label_lr plant_lr = args.plant_lr extractor_lr = args.extractor_lr domain_adapt_lr = args.domain_adapt_lr device = 'cuda' if torch.cuda.is_available() else 'cpu' curr_experiment = experiments_info[args.experiment] modalities = get_experiment_modalities_params(curr_experiment, args.lwir_skip, args.lwir_max_len, args.vir_max_len, args.color_max_len) used_modalities = get_used_modalities(modalities, args.excluded_modalities) if args.experiment_path is None: experiment_path = args.experiment else: experiment_path = args.experiment_path end_date = curr_experiment.end_date if args.num_days is not None: end_date = curr_experiment.start_date + timedelta(days=args.num_days - 1) dataset = Modalities(experiment_path, args.experiment, split_cycle=args.split_cycle, start_date=curr_experiment.start_date, end_date=end_date, **used_modalities) train_set, test_set = ModalitiesSubset.random_split( dataset, args.train_ratio) train_loader = data.DataLoader(train_set, batch_size=batch_size, num_workers=2, shuffle=True) feat_extractor_params = dict() for mod in used_modalities.keys(): num_levels, kernel_size = get_levels_kernel( dataset.modalities[mod].max_len) feat_extractor_params[mod] = { 'num_levels': num_levels, 'kernel_size': kernel_size } feat_ext = FeatureExtractor(**feat_extractor_params).to(device) label_cls = nn.Sequential(nn.ReLU(), nn.Linear(512, len( classes[args.experiment]))).to(device) plant_cls = nn.Sequential(nn.ReLU(), nn.Linear(512, train_set.num_plants)).to(device) criterion = nn.CrossEntropyLoss(reduction='sum').to(device) label_opt = optim.SGD(label_cls.parameters(), lr=label_lr, weight_decay=1e-3) plant_opt = optim.SGD(plant_cls.parameters(), lr=plant_lr, weight_decay=1e-3) ext_opt = optim.SGD(feat_ext.parameters(), lr=extractor_lr, weight_decay=1e-3) best_loss = float('inf') test_config = TestConfig(args.use_checkpoints, checkpoint_name, epochs, batch_size, domain_adapt_lr, device, dataset, train_set, test_set, train_loader, feat_ext, label_cls, plant_cls, criterion, label_opt, plant_opt, ext_opt, best_loss, args.loss_delta, args.return_epochs) if args.load_checkpoint: restore_checkpoint(test_config) train_loop(test_config) if args.use_checkpoints: restore_checkpoint(test_config)
list_total_user_labels.append(user_labels) # total_user_data = [] # total_user_labels = [] # for subject_index in range(11): # total_user_data.extend(list_total_user_data[subject_index]) # total_user_labels.extend(list_total_user_labels[subject_index]) total_user_data = np.array(list_total_user_data, dtype=np.float32) # <np.ndarray> (11, 180, 400, 8) total_user_labels = np.array(list_total_user_labels, dtype=np.int64) # <np.ndarray> (11, 180) # ----------------------------------------- Init Network -------------------------------------------------------- # feature_extractor = FeatureExtractor().cuda() domain_classifier = DomainClassifier().cuda() label_predictor = LabelPredictor().cuda() feature_extractor.load_state_dict( torch.load(r'saved_model\feature_extractor_CE_8_subjects.pkl')) domain_classifier.load_state_dict( torch.load(r'saved_model\domain_classifier_CE_8_subjects.pkl')) label_predictor.load_state_dict( torch.load(r'saved_model\label_predictor_CE_8_subjects.pkl')) # ------------------------------------------ Testing Stage -------------------------------------------------------- # window_size = 52 stride = 1 max_fit = 30
def main(args): np.random.seed(0) torch.manual_seed(0) with open('config.yaml', 'r') as file: stream = file.read() config_dict = yaml.safe_load(stream) config = mapper(**config_dict) disc_model = Discriminator(input_shape=(config.data.channels, config.data.hr_height, config.data.hr_width)) gen_model = GeneratorResNet() feature_extractor_model = FeatureExtractor() plt.ion() if config.distributed: disc_model.to(device) disc_model = nn.parallel.DistributedDataParallel(disc_model) gen_model.to(device) gen_model = nn.parallel.DistributedDataParallel(gen_model) feature_extractor_model.to(device) feature_extractor_model = nn.parallel.DistributedDataParallel( feature_extractor_model) elif config.gpu: # disc_model = nn.DataParallel(disc_model).to(device) # gen_model = nn.DataParallel(gen_model).to(device) # feature_extractor_model = nn.DataParallel(feature_extractor_model).to(device) disc_model = disc_model.to(device) gen_model = gen_model.to(device) feature_extractor_model = feature_extractor_model.to(device) else: return train_dataset = ImageDataset(config.data.path, hr_shape=(config.data.hr_height, config.data.hr_width), lr_shape=(config.data.lr_height, config.data.lr_width)) test_dataset = ImageDataset(config.data.path, hr_shape=(config.data.hr_height, config.data.hr_width), lr_shape=(config.data.lr_height, config.data.lr_width)) if config.distributed: train_sampler = torch.utils.data.distributed.DistributedSampler( train_dataset) else: train_sampler = None train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=config.data.batch_size, shuffle=config.data.shuffle, num_workers=config.data.workers, pin_memory=config.data.pin_memory, sampler=train_sampler) val_loader = torch.utils.data.DataLoader(test_dataset, batch_size=config.data.batch_size, shuffle=config.data.shuffle, num_workers=config.data.workers, pin_memory=config.data.pin_memory) if args.train: # trainer settings trainer = GANTrainer(config.train, train_loader, (disc_model, gen_model, feature_extractor_model)) criterion = nn.MSELoss().to(device) disc_optimizer = torch.optim.Adam(disc_model.parameters(), config.train.hyperparameters.lr) gen_optimizer = torch.optim.Adam(gen_model.parameters(), config.train.hyperparameters.lr) fe_optimizer = torch.optim.Adam(feature_extractor_model.parameters(), config.train.hyperparameters.lr) trainer.setCriterion(criterion) trainer.setDiscOptimizer(disc_optimizer) trainer.setGenOptimizer(gen_optimizer) trainer.setFEOptimizer(fe_optimizer) # evaluator settings evaluator = GANEvaluator( config.evaluate, val_loader, (disc_model, gen_model, feature_extractor_model)) # optimizer = torch.optim.Adam(disc_model.parameters(), lr=config.evaluate.hyperparameters.lr, # weight_decay=config.evaluate.hyperparameters.weight_decay) evaluator.setCriterion(criterion) if args.test: pass # Turn on benchmark if the input sizes don't vary # It is used to find best way to run models on your machine cudnn.benchmark = True start_epoch = 0 best_precision = 0 # optionally resume from a checkpoint if config.train.resume: [start_epoch, best_precision] = trainer.load_saved_checkpoint(checkpoint=None) # change value to test.hyperparameters on testing for epoch in range(start_epoch, config.train.hyperparameters.total_epochs): if config.distributed: train_sampler.set_epoch(epoch) if args.train: trainer.adjust_learning_rate(epoch) trainer.train(epoch) prec1 = evaluator.evaluate(epoch) if args.test: pass # remember best prec@1 and save checkpoint if args.train: is_best = prec1 > best_precision best_precision = max(prec1, best_precision) trainer.save_checkpoint( { 'epoch': epoch + 1, 'state_dict': disc_model.state_dict(), 'best_precision': best_precision, 'optimizer': optimizer.state_dict(), }, is_best, checkpoint=None)