def __init__(self, args, vocab_len): super(BFM, self).__init__() self.encoder = EncoderCNN(args.embed_size).eval().cpu() self.encoder.load_state_dict(torch.load('encoder.ckpt', map_location=torch.device('cpu'))) self.decoder = DecoderRNN(args.embed_size, args.hidden_size, vocab_len, args.num_layers).eval().cpu() self.decoder.forward = self.decoder.sample self.decoder.load_state_dict(torch.load('decoder.ckpt', map_location=torch.device('cpu')))
def __init__(self): print("Defining I.A") # Device configuration self.device = torch.device('cpu') #vars embed_size = 256 hidden_size = 512 num_layers = 1 encoder_path = 'models/encoder-5-3000.pkl' decoder_path = 'models/decoder-5-3000.pkl' vocab_path = 'data/vocab.pkl' # Image preprocessing self.transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))]) with open(vocab_path, 'rb') as f: self.vocab = pickle.load(f) print("Building Model") # Build models self.encoder = EncoderCNN(embed_size).eval() # eval mode (batchnorm uses moving mean/variance) self.decoder = DecoderRNN(embed_size, hidden_size, len(self.vocab), num_layers) self.encoder = self.encoder.to(self.device) self.decoder = self.decoder.to(self.device) print("loading checkpoint") # Load the trained model parameters self.encoder.load_state_dict(torch.load(encoder_path)) self.decoder.load_state_dict(torch.load(decoder_path))
def main(args): # Image preprocessing, normalization for the pretrained resnet transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))]) # Load vocabulary wrapper with open(args.vocab_path, 'rb') as f: vocab = pickle.load(f) # Build models encoder = EncoderCNN(args.embed_size).eval() # eval mode (batchnorm uses moving mean/variance) decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), 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)) # Build data loader data_loader = get_loader(args.image_dir, args.caption_path, vocab, transform, args.batch_size, shuffle=True, num_workers=args.num_workers) total_step = len(data_loader) # List to score the BLEU scores bleu_scores = [] for i, (images, captions, lengths) in enumerate(data_loader): # Set mini-batch dataset images = images.to(device) # captions = captions.to(device) # Generate an caption from the image feature = encoder(images) sampled_ids = decoder.sample(feature) sampled_ids = sampled_ids[0].cpu().numpy() # Convert word_ids to words sampled_caption = [] for word_id in sampled_ids: word = vocab.idx2word[word_id] sampled_caption.append(word) if word == '<end>': break sentence = ' '.join(sampled_caption) score = sentence_bleu(captions, sentence, args.bleu_weights) bleu_scores.append(score) # Print log info if i % args.log_step == 0: print('Finish [{}/{}], Current BLEU Score: {:.4f}' .format(i, total_step, np.mean(bleu_scores))) np.save('test_results.npy', [bleu_scores, np.mean(bleu_scores)])
def test(args): transform = transforms.Compose([ transforms.ToTensor(), ]) with open(args.vocab_path, 'rb') as f: vocab = pickle.load(f) encoder = EncoderCNN(args.embed_size).eval() decoder = DecoderRNN(args.embed_size, len(vocab), args.hidden_size, args.num_layers) # 加载训练好的模型的参数 encoder.load_state_dict(torch.load(args.encoder_path, map_location='cpu')) decoder.load_state_dict(torch.load(args.decoder_path, map_location='cpu')) image = load_img(args.img_path, transform) feature = encoder(image) sampled_ids = decoder.sample(feature) sampled_ids = sampled_ids[0].cpu().numpy() sampled_caption = [] for word_id in sampled_ids: word = vocab.idx2word[word_id] sampled_caption.append(word) if word == '<end>': break sentence = ' '.join(sampled_caption) print(sentence) image = Image.open(args.img_path) plt.imshow(np.asarray(image)) plt.show()
def getCaption(self, imgs, output_path='', vocab_path='data/vocab.pkl', decoder_path='models/decoder-5-3000.pkl', encoder_path='models/encoder-5-3000.pkl', embed_size=256, hidden_size=512, num_layers=1): if (output_path == ''): output_path = self.DEFAULT_OUTPUT_PATH device = self.device transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) # Load vocabulary wrapper with open(vocab_path, 'rb') as f: vocab = pickle.load(f) # Build models encoder = EncoderCNN(embed_size).eval( ) # eval mode (batchnorm uses moving mean/variance) decoder = DecoderRNN(embed_size, hidden_size, len(vocab), num_layers) encoder = encoder.to(device) decoder = decoder.to(device) # Load the trained model parameters encoder.load_state_dict(torch.load(encoder_path)) decoder.load_state_dict(torch.load(decoder_path)) CAPTIONS = [] for img in imgs: # Prepare an image image = self.load_image(img, transform=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 = vocab.idx2word[word_id] sampled_caption.append(word) if word == '<end>': break sentence = ' '.join(sampled_caption) # Print out the image and the generated caption CAPTIONS.append(self.prune_caption(sentence)) json_captions = self.writeJSON(imgs, CAPTIONS, output_path=output_path) return json_captions
def __init__(self): self.transform = transforms.Compose([ transforms.Resize(256), # smaller edge of image resized to 256 transforms.CenterCrop(224), # get 224x224 crop from the center transforms.ToTensor(), # convert the PIL Image to a tensor transforms.Normalize((0.485, 0.456, 0.406), # normalize image for pre-trained model (0.229, 0.224, 0.225))]) # Load cherckpoint with best model self.checkpoint = torch.load(os.path.join('./models', 'best-model.pkl'), 'cpu') # Specify values for embed_size and hidden_size - we use the same values as in training step self.embed_size = 512 self.hidden_size = 512 # Get the vocabulary and its size self.vocab = Vocabulary(None, './vocab.pkl', "<start>", "<end>", "<unk>", "<pad>", "", "", True) self.vocab_size = len(self.vocab) # Initialize the encoder and decoder, and set each to inference mode self.encoder = EncoderCNN(self.embed_size) self.encoder.eval() self.decoder = DecoderRNN(self.embed_size, self.hidden_size, self.vocab_size) self.decoder.eval() # Load the pre-trained weights self.encoder.load_state_dict(self.checkpoint['encoder']) self.decoder.load_state_dict(self.checkpoint['decoder'])
def main(): # Load vocabulary wrapper. with open(vocab_path) as f: vocab = pickle.load(f) encoder = EncoderCNN(4096, embed_dim) encoder.load_state_dict(torch.load('searchimage.pkl')) for p in encoder.parameters(): p.requires_grad = False word_encoder = EncoderRNN(embed_dim, embed_dim, len(vocab), num_layers_rnn) word_encoder.load_state_dict(torch.load('searchword.pkl')) if torch.cuda.is_available(): encoder.cuda() word_encoder.cuda() # Loss and Optimizer criterion = nn.MSELoss() params = list( word_encoder.parameters()) # + list(encoder.linear.parameters()) optimizer = torch.optim.Adam(params, lr=2e-6, weight_decay=0.001) #load data with open(image_data_file) as f: image_data = pickle.load(f) image_features = si.loadmat(image_feature_file) img_features = image_features['fc7'][0] img_features = np.concatenate(img_features) print 'here' iteration = 0 for i in range(10): # epoch use_caption = i % 5 print 'Epoch', i losses = [] for x, y in make_mini_batch(img_features, image_data, use_caption=use_caption): encoder.zero_grad() word_encoder.zero_grad() word_padding, lengths = make_word_padding(y, vocab) x = Variable(torch.from_numpy(x).cuda()) word_index = Variable(torch.from_numpy(word_padding).cuda()) features = encoder(x) outputs = word_encoder(word_index, lengths) loss = torch.mean((features - outputs).pow(2)) loss.backward() optimizer.step() losses.append(loss.data[0]) if iteration % 100 == 0: print 'loss', sum(losses) / float(len(losses)) losses = [] iteration += 1 torch.save(word_encoder.state_dict(), 'searchword.pkl') torch.save(encoder.state_dict(), 'searchimage.pkl')
def run_inference(image_path, encoder_path, decoder_path, vocab_path, embed_size=256, hidden_size=512, num_layers=1): print(f'sample.py running ... ') # Image preprocessing transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) # Load vocabulary wrapper with open(vocab_path, 'rb') as f: print("using " + vocab_path) vocab = pickle.load(f) # Build models encoder = EncoderCNN( embed_size).eval() # eval mode (batchnorm uses moving mean/variance) decoder = DecoderRNN(embed_size, hidden_size, len(vocab), num_layers) encoder = encoder.to(device) decoder = decoder.to(device) # Load the trained model parameters encoder.load_state_dict( torch.load(encoder_path, map_location=torch.device('cpu'))) decoder.load_state_dict( torch.load(decoder_path, map_location=torch.device('cpu'))) # Prepare an image image = load_image(image_path, 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 = vocab.idx2word[word_id] print(word) sampled_caption.append(word) if word == '<end>': break sentence = ' '.join(sampled_caption).replace('<start>', '') sentence = sentence.replace('<end>', '') sentence = sentence.replace('_', ' ') # Print out the image and the generated caption print(sentence) print(f'debug: chay xong roi ne') return sentence.strip().capitalize()
def inference_coco(encoder_file: str, decoder_file: str, embed_size: int, hidden_size: int, from_cpu: bool) -> None: """ Displays an original image from coco test dataset and prints its associated caption. encoder_file: Name of the encoder to load. decoder_file: Name of the decoder to load. embed_size: Word embedding size for the encoder. hidden_size: Hidden layer of the LSTM size. from_cpu: Whether the model has been saved on CPU. """ # Define transform transform_test = transforms.Compose([ transforms.Resize(256), # smaller edge of image resized to 256 transforms.RandomCrop(224), # get 224x224 crop from random location transforms.ToTensor(), # convert the PIL Image to a tensor transforms.Normalize( (0.485, 0.456, 0.406), # normalize image for pre-trained model (0.229, 0.224, 0.225)) ]) # Device to use fo inference device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Create the data loader. data_loader = get_loader(transform=transform_test, mode='test') # Obtain sample image _, image = next(iter(data_loader)) # The size of the vocabulary. vocab_size = len(data_loader.dataset.vocab) # Initialize the encoder and decoder, and set each to inference mode. encoder = EncoderCNN(embed_size) encoder.eval() decoder = DecoderRNN(embed_size, hidden_size, vocab_size) decoder.eval() # Load the trained weights. if from_cpu: encoder.load_state_dict( torch.load(os.path.join('./models', encoder_file), map_location='cpu')) decoder.load_state_dict( torch.load(os.path.join('./models', decoder_file), map_location='cpu')) else: encoder.load_state_dict( torch.load(os.path.join('./models', encoder_file))) decoder.load_state_dict( torch.load(os.path.join('./models', decoder_file))) # Move models to GPU if CUDA is available. encoder.to(device) decoder.to(device) get_prediction(encoder, decoder, data_loader, device)
def main(args): # Create model directory if not os.path.exists(args.model_path): os.makedirs(args.model_path) # Build the models encoder = EncoderCNN(args.embed_size).to(device) decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), args.num_layers).to(device) # Loss and optimizer criterion = nn.CrossEntropyLoss() params = list(decoder.parameters()) + list( encoder.linear.parameters()) + list(encoder.bn.parameters()) optimizer = torch.optim.Adam(params, lr=args.learning_rate) # Train the models total_step = len(data_loader) # For each TSP problem for epoch in range(args.num_epochs): for i, (images, captions, lengths) in enumerate(data_loader): # Set mini-batch dataset images = images.to(device) captions = captions.to(device) targets = pack_padded_sequence(captions, lengths, batch_first=True)[0] # Forward, backward and optimize features = encoder(images) outputs = decoder(features, captions, lengths) loss = criterion(outputs, targets) decoder.zero_grad() encoder.zero_grad() loss.backward() optimizer.step() # Print log info if i % args.log_step == 0: print( 'Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}, Perplexity: {:5.4f}' .format(epoch, args.num_epochs, i, total_step, loss.item(), np.exp(loss.item()))) # Save the model checkpoints if (i + 1) % args.save_step == 0: torch.save( decoder.state_dict(), os.path.join(args.model_path, 'decoder-{}-{}.ckpt'.format(epoch + 1, i + 1))) torch.save( encoder.state_dict(), os.path.join(args.model_path, 'encoder-{}-{}.ckpt'.format(epoch + 1, i + 1)))
def main(args): # Image preprocessing prediction = [] transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))]) # Load vocabulary wrapper with open(args.vocab_path, 'rb') as f: vocab = pickle.load(f) # Build models encoder = EncoderCNN(args.embed_size).eval() # eval mode (batchnorm uses moving mean/variance) decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), 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)) dirname='' fnames = listdir(os.getcwd) #with open(dirname) for fname in fnames: #print(fname) # Prepare an image image = load_image(''+fname, 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 = vocab.idx2word[word_id] sampled_caption.append(word) if word == '<end>': break sentence = ' '.join(sampled_caption) current_pred = [fname, sentence] predictions.append(current_pred) # Print out the image and the generated caption print(fname) print (sentence) #image = Image.open(args.image) #plt.imshow(np.asarray(image)) df = pd.DataFrame(predictions, columns=['File Name', 'Caption']) df.to_excel('output.xls')
class Neuraltalk2: def __init__(self): print("Defining I.A") # Device configuration self.device = torch.device('cpu') #vars embed_size = 256 hidden_size = 512 num_layers = 1 encoder_path = 'models/encoder-5-3000.pkl' decoder_path = 'models/decoder-5-3000.pkl' vocab_path = 'data/vocab.pkl' # Image preprocessing self.transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))]) with open(vocab_path, 'rb') as f: self.vocab = pickle.load(f) print("Building Model") # Build models self.encoder = EncoderCNN(embed_size).eval() # eval mode (batchnorm uses moving mean/variance) self.decoder = DecoderRNN(embed_size, hidden_size, len(self.vocab), num_layers) self.encoder = self.encoder.to(self.device) self.decoder = self.decoder.to(self.device) print("loading checkpoint") # Load the trained model parameters self.encoder.load_state_dict(torch.load(encoder_path)) self.decoder.load_state_dict(torch.load(decoder_path)) def eval_image(self, image_path): # Prepare an image image = load_image(image_path, self.transform) image_tensor = image.to(self.device) # Generate an caption from the image feature = self.encoder(image_tensor) sampled_ids = self.decoder.sample(feature) sampled_ids = sampled_ids[0].cpu().numpy() # Convert word_ids to words sampled_caption = [] for word_id in sampled_ids: word = self.vocab.idx2word[word_id] if word == '<end>': break if word == '<start>': continue sampled_caption.append(word) sentence = ' '.join(sampled_caption) return sentence
def predict(self, args): print('predict..start') device = torch.device('cpu') # Image preprocessing transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) # Load vocabulary wrapper # Load vocabulary wrapper #vocab = Vocabulary() with open(args.vocab_path, 'rb') as f: vocab = pickle.load(f) # Build models encoder = EncoderCNN(args.embed_size).eval( ) # eval mode (batchnorm uses moving mean/variance) decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), 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, map_location=device)) decoder.load_state_dict( torch.load(args.decoder_path, map_location=device)) # Prepare an image image = self.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 = vocab.idx2word[word_id] sampled_caption.append(word) if word == '<end>': break sentence = ' '.join(sampled_caption) # Print out the image and the generated caption print(sentence) return sentence
def main(cfg): # print(cfg.pretty()) transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(cfg.image.mean, cfg.image.std) ]) print(hydra.utils.to_absolute_path(cfg.train.vocab_path)) with open(hydra.utils.to_absolute_path(cfg.train.vocab_path), 'rb') as f: vocab = pickle.load(f) # モデルの構築 encoder = EncoderCNN(cfg.train.embed_size).eval() decoder = DecoderRNN(cfg.train.embed_size, cfg.train.hidden_size, len(vocab), cfg.train.num_layers) encoder = encoder.to(device) decoder = decoder.to(device) # 学習済みモデルのパラメータを読み込む encoder.load_state_dict( torch.load(hydra.utils.to_absolute_path(cfg.train.encoder_path))) decoder.load_state_dict( torch.load(hydra.utils.to_absolute_path(cfg.train.decoder_path))) with open(json_dir, encoding='utf-8') as f: data = json.loads(f.read()) for key in data['images']: img_file_name = key['file_name'] img_file_path = base_dir + '/data/val2014/' + img_file_name # 画像の準備 image = load_image(hydra.utils.to_absolute_path(img_file_path), transform) image_tensor = image.to(device) # 入力した画像からキャプションを生成する feature = encoder(image_tensor) sampled_ids = decoder.sample(feature) sampled_ids = sampled_ids[0].cpu().numpy() # word_idsをwordに変換する sampled_caption = [] for word_id in sampled_ids: word = vocab.idx2word[word_id] sampled_caption.append(word) if word == '<end>': break # <start>,<end>,"."を取り除いて処理を行う sentence = ' '.join(sampled_caption[1:-2]) print(sentence)
def main2(image, encoder_path='models/encoder-5-3000.pkl', decoder_path='models/decoder-5-3000.pkl', vocab_path="data/vocab.pkl", embed_size=256, hidden_size=512, num_layers=1): # Image preprocessing transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) # Load vocabulary wrapper with open(vocab_path, 'rb') as f: vocab = pickle.load(f) # Build models encoder = EncoderCNN( embed_size).eval() # eval mode (batchnorm uses moving mean/variance) decoder = DecoderRNN(embed_size, hidden_size, len(vocab), num_layers) encoder = encoder.to(device) decoder = decoder.to(device) # Load the trained model parameters encoder.load_state_dict(torch.load(encoder_path)) decoder.load_state_dict(torch.load(decoder_path)) # Prepare an image image = load_image(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 = vocab.idx2word[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)) return sentence
class BFM(nn.Module): def __init__(self, args, vocab_len): super(BFM, self).__init__() self.encoder = EncoderCNN(args.embed_size).eval().cpu() self.encoder.load_state_dict(torch.load('encoder.ckpt', map_location=torch.device('cpu'))) self.decoder = DecoderRNN(args.embed_size, args.hidden_size, vocab_len, args.num_layers).eval().cpu() self.decoder.forward = self.decoder.sample self.decoder.load_state_dict(torch.load('decoder.ckpt', map_location=torch.device('cpu'))) def forward(self, image): feature = self.encoder(image) sampled_ids = self. decoder(feature) return sampled_ids
def main(args): # Image preprocessing transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) # Load vocabulary wrapper with open(args.vocab_path, 'rb') as f: vocab = pickle.load(f) # Build models encoder = EncoderCNN(args.embed_size).eval( ) # eval mode (batchnorm uses moving mean/variance) decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), 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 = vocab.idx2word[word_id] sampled_caption.append(word) if word == '<end>': break sampled_caption.remove('<end>') sampled_caption.remove('<start>') sentence = ' '.join(sampled_caption) # Print out the image and the generated caption print(sentence) f = open("demofile3.txt", "w") f.truncate(0) f.write(sentence) f.close() image = Image.open(args.image) plt.imshow(np.asarray(image))
def main(): # TODO parse arguments parser = argparse.ArgumentParser() # DATA related parser.add_argument('--data_path', type=str, help='Path of training data') parser.add_argument('--max_length', type=int, default=250, help='Pad all data to this size') # Model Specs parser.add_argument('--embed_size', type=int, default=300, help='Dimension of word embedding') parser.add_argument('--filter_size', type=int, default=5, help='size of filter') parser.add_argument('--stride', type=int, default=2, help='stride size for each layer') parser.add_argument('--filter_nums', type=str, default='300,600', help='filer number for each convolution layer') parser.add_argument('--hidden_size', type=int, default=500, help='size of hidden state in the middle') # Optimization Specs parser.add_argument('--batch_size', type=int, default=128, help='Minibatch during training') parser.add_argument('--optim', type=str, default='SGD', help='Optimization method') parser.add_argument('--learning_rate', type=float, default=0.0005, help='Learning rate of model') parser.add_argument('--decay', type=float, default=0.5, help='Decay rate of learning rate') parser.add_argument('--start_decay', type=int, default=5, help='Start Epoch of decay learning rate') args = parser.parse_args() args.filter_nums = [int(i) for i in args.filter_nums.split(',')] args.filter_nums.append(args.hidden_size) final_filter_size = args.max_length for num in args.filter_nums[:-1]: final_filter_size = math.floor( (final_filter_size - args.filter_size)/args.stride + 1 ) # TODO build model encoder = EncoderCNN(args.embed_size, args.filter_size, args.stride, args.filter_nums, final_filter_size) decoder = DecoderCNN(args.embed_size, args.filter_size, args.stride, args.filter_nums, final_filter_size) if torch.cuda.is_available(): encoder.cuda() decoder.cuda() creterion = # TODO Load training data with open(args.data_path, 'r') as f: data = f.readlines()
def main(args): with open('data/vocab.pkl', 'rb') as f: vocab = pickle.load(f) encoder = EncoderCNN(256) encoder.eval() # evaluation mode (BN uses moving mean/variance) decoder = DecoderRNN(256, 512, len(vocab), 1) if torch.cuda.is_available(): encoder.cuda() decoder.cuda() # Load the trained model parameters encoder.load_state_dict(torch.load(args.encoder)) decoder.load_state_dict(torch.load(args.decoder)) measurement_score = test(encoder, decoder, vocab, args.num_samples, args.num_hints, args.debug, args.c_step, args.no_avg) if args.msm == "co": scores = cocoEval() scores_u = cocoEval(res='data/captions_val2014_results_u.json') print(scores) print(scores_u) with open(args.filepath, 'w+') as f: pickle.dump((scores, scores_u), f)
def main(args): # Build models encoder = EncoderCNN(args.embed_size).eval( ) # eval mode (batchnorm uses moving mean/variance) decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), 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)) #path for dumbing output of encoder model decoder.load_state_dict(torch.load( args.decoder_path)) #path for dumbing output of decoder model test_data = generate_training_data(5) textures_test = generate_textures(test_data) transforms_test = generate_transforms(test_data) for i in range(len(textures_test)): plt.imsave('predictions/texture4_0%i.png' % i, textures_test[i], cmap="gray") print(transforms_test) predicted_progs = [] for texture in textures_test: texture = torch.tensor(texture, device=device) texture = texture.unsqueeze(0) texture = texture.unsqueeze( 0) #for EncoderCNN ought to unsqueeze twice feature = encoder(texture) sampled_seq = decoder.sample(feature) sampled_seq = sampled_seq[0].cpu().numpy( ) # (1, max_seq_length) -> (max_seq_length) # Convert sampled sequence of transforms to words prog = [] for int_word in sampled_seq: word = int_to_word(int_word) prog.append(word) if word == '<end>': break trans_seq = '-->'.join(prog) predicted_progs.append([trans_seq]) # Print out the sequence of generated transform sequences print(predicted_progs)
def main(args): # Val images folder filepath = '/scratch/ys2542/pytorch-tutorial/tutorials/03-advanced/image_captioning/data/resizedval2014' onlyfiles = [fl for fl in listdir(filepath) if isfile(join(filepath, fl))] # image preprocessing transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) # load vocabulary wrapper pickle file with open(args.vocab_path, 'rb') as f: vocab = pickle.load(f) encoder = EncoderCNN(args.embed_size) # build encoder encoder.eval() # evaluation mode by moving mean and variance decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), args.num_layers) # build decoder # load the trained CNN and RNN parameters encoder.load_state_dict(torch.load(args.encoder_path)) decoder.load_state_dict(torch.load(args.decoder_path)) # Load all images in val folder for i in onlyfiles: badsize = 0 # count the unload images args_image = filepath + '/' # val folder path with image names args_image = args_image + i # transform image and wrap it to tensor image = load_image(args_image, transform) image_tensor = to_var(image, volatile=True) if torch.cuda.is_available(): # load GPU encoder.cuda() decoder.cuda() # generate caption from image try: feature = encoder(image_tensor) sampled_ids = decoder.sample(feature) sampled_ids = sampled_ids.cpu().data.numpy() # decode word_ids to words sampled_caption = [] for word_id in sampled_ids: word = vocab.idx2word[word_id] sampled_caption.append(word) if word == '<end>': break sentence = ' '.join(sampled_caption) # print out image and generated caption without start and end print('beam_size_1' + '\t' + i + '\t' + sentence[8:-8]) except: badsize = badsize + 1 # count some wrong images
def main(image): # Configuration for hyper-parameters config = Config() # Image Preprocessing transform = config.test_transform # Load vocabulary with open(os.path.join(config.vocab_path, 'vocab.pkl'), 'rb') as f: vocab = pickle.load(f) # Build Models encoder = EncoderCNN(config.embed_size) encoder.eval() # evaluation mode (BN uses moving mean/variance) decoder = DecoderRNN(config.embed_size, config.hidden_size, len(vocab), config.num_layers) # Load the trained model parameters encoder.load_state_dict( torch.load( os.path.join(config.teacher_cnn_path, config.trained_encoder))) decoder.load_state_dict( torch.load( os.path.join(config.teacher_lstm_path, config.trained_decoder))) # Prepare Image image = Image.open(image) image_tensor = Variable(transform(image).unsqueeze(0)) # Set initial states state = (Variable(torch.zeros(config.num_layers, 1, config.hidden_size)), Variable(torch.zeros(config.num_layers, 1, config.hidden_size))) # If use gpu if torch.cuda.is_available(): encoder.cuda() decoder.cuda() state = [s.cuda() for s in state] image_tensor = image_tensor.cuda() # Generate caption from image feature = encoder(image_tensor) sampled_ids = decoder.sample(feature, state) sampled_ids = sampled_ids.cpu().data.numpy() # Decode word_ids to words sampled_caption = [] for word_id in sampled_ids: word = vocab.idx2word[word_id] sampled_caption.append(word) if word_id == 96: sampled_caption.append('<end>') break if word == '<end>': break sentence = ' '.join(sampled_caption) # Print out image and generated caption. print(sentence) return sentence
def main(args): # Image preprocessing transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) # Load vocabulary wrapper with open(args.vocab_path, 'rb') as f: vocab = pickle.load(f) # Build Models encoder = EncoderCNN(args.embed_size) encoder.eval() # evaluation mode (BN uses moving mean/variance) decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), args.num_layers) # Load the trained model parameters encoder.load_state_dict(torch.load(args.encoder_path)) decoder.load_state_dict(torch.load(args.decoder_path)) # Prepare Image image = load_image(args.image, transform) image_tensor = to_var(image, volatile=True) # If use gpu if torch.cuda.is_available(): encoder.cuda() decoder.cuda() # Generate caption from image feature = encoder(image_tensor) sampled_ids = decoder.sample(feature) sampled_ids = sampled_ids.cpu().data.numpy() # Decode word_ids to words sampled_caption = [] for word_id in sampled_ids: word = vocab.idx2word[word_id] sampled_caption.append(word) if word == '<end>': break sentence = ' '.join(sampled_caption) # Print out image and generated caption. sentence = sentence.replace('<start> ', '').replace(' <end>', '').replace('.', '').strip() translator = Translator() sentence_indo = translator.translate(sentence, dest='id').text print('This is an image of: ' + sentence_indo) tts = gTTS(sentence_indo, 'id') tts.save('result.mp3') playsound('result.mp3') image = Image.open(args.image) plt.imshow(np.asarray(image)) plt.show()
def main(args): # Image preprocessing transform = transforms.Compose([ transforms.Scale(args.crop_size), transforms.CenterCrop(args.crop_size), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) # Load vocabulary wrapper with open(args.vocab_path, 'rb') as f: vocab = pickle.load(f) alexnet = models.alexnet(pretrained=True) alexnet2 = AlexNet2(alexnet) # Build Models encoder = EncoderCNN(4096, args.embed_size) encoder.eval() # evaluation mode (BN uses moving mean/variance) decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), args.num_layers) # Load the trained model parameters encoder.load_state_dict(torch.load(args.encoder_path)) decoder.load_state_dict(torch.load(args.decoder_path)) # Prepare Image image = Image.open(args.image) image_tensor = Variable(transform(image).unsqueeze(0)) # Set initial states state = (Variable(torch.zeros(args.num_layers, 1, args.hidden_size)), Variable(torch.zeros(args.num_layers, 1, args.hidden_size))) # If use gpu if torch.cuda.is_available(): encoder.cuda() decoder.cuda() alexnet2.cuda() state = [s.cuda() for s in state] image_tensor = image_tensor.cuda() # Generate caption from image alexnet2(image_tensor) feature = encoder(alexnet2.fc7_value) sampled_ids = decoder.sample(feature, state) sampled_ids = sampled_ids.cpu().data.numpy() # Decode word_ids to words sampled_caption = [] for word_id in sampled_ids: word = vocab.idx2word[word_id] sampled_caption.append(word) if word == '<end>': break sentence = ' '.join(sampled_caption) # Print out image and generated caption. print(sentence)
def main(args): # Image preprocessing transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))]) # Load vocabulary wrapper with open(args.vocab_path, 'rb') as f: vocab = pickle.load(f) # Build models encoder = EncoderCNN(args.embed_size).eval() # eval mode (batchnorm uses moving mean/variance) decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), 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) ''' data_loader, _ = get_loader(transforms=False) inp, targets = next(iter(data_loader)) audio = inp_transform(inp) audio = audio.to(device) # Generate an caption from the image feature = encoder(audio) 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 = vocab.idx2word[word_id] sampled_caption.append(word) if word == '<end>': break sentence = ' '.join(sampled_caption) # Print out the image and the generated caption print("Logits : {}\nTarget : {}".format(sentence, targets)) '''
def get_encoder_decoder(vocab): """ Given the arguments, returns the correct combination of CNN/RNN/GAN encoders and decoders. """ if args.pretrain_rnn: encoder = EncoderRNN(len(vocab), args.embed_size, args.encoder_rnn_hidden_size, num_layers=args.num_layers).to(device) elif args.gan_embedding: gan = torch.load('DCGAN_embed_2.tch').to(device) encoder = gan.discriminator elif args.progan_embedding: pro_gan = pg.ProGAN(depth=7, latent_size=256, device=torch.device('cuda')) pro_gan.dis.load_state_dict(torch.load('progan_weights/GAN_DIS_6.pth')) # pro_gan.dis_optim.load_state_dict(torch.load('progan_weights/GAN_DIS_OPTIM_6.pth')) pro_gan.gen.load_state_dict(torch.load('progan_weights/GAN_GEN_6.pth')) # pro_gan.gen_optim.load_state_dict(torch.load('progan_weights/GAN_GEN_OPTIM_6.pth')) pro_gan.gen_shadow.load_state_dict( torch.load('progan_weights/GAN_GEN_SHADOW_6.pth')) print("Loaded proGAN weights.", flush=True) encoder = pro_gan.dis.to(device) else: encoder = EncoderCNN(args.embed_size).to(device) decoder = DecoderRNNOld(args.embed_size, args.decoder_rnn_hidden_size, len(vocab), args.num_layers, vocab, device=device).to(device) return encoder, decoder
def main(args): # Image preprocessing transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) # Load vocabulary wrapper with open(args.vocab_path, 'rb') as f: vocab = pickle.load(f) # Build Models encoder = EncoderCNN(args.embed_size) encoder.eval() # evaluation mode (BN uses moving mean/variance) decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), args.num_layers) # Load the trained model parameters encoder.load_state_dict(torch.load(args.encoder_path)) decoder.load_state_dict(torch.load(args.decoder_path)) # Prepare Image #image = load_image(args.image, transform) #image_tensor = to_var(image, volatile=True) # If use gpu if torch.cuda.is_available(): encoder.cuda() decoder.cuda() data = [] try: img_path = args.image # Prepare Image image = load_image(img_path, transform) image_tensor = to_var(image, volatile=True) # Generate caption from image feature = encoder(image_tensor) #pdb.set_trace() sampled_ids = decoder.sample(feature) sampled_ids = sampled_ids.cpu().data.numpy() # Decode word_ids to words sampled_caption = [] for word_id in sampled_ids: word = vocab.idx2word[word_id] if word == '<start>': continue if word == '<end>': break sampled_caption.append(word) sentence = ' '.join(sampled_caption) # Print out image and generated caption. print(sentence) data.append({'key': img_path.split('/')[-1], 'sentence': sentence}) except: print(img_path)
def main(args): # Image preprocessing transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))]) # Load vocabulary wrapper with open(args.vocab_path, 'rb') as f: vocab = pickle.load(f) # Build Models encoder = EncoderCNN(args.embed_size) encoder.eval() # evaluation mode (BN uses moving mean/variance) decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), args.num_layers) # Load the trained model parameters encoder.load_state_dict(torch.load(args.encoder_path)) decoder.load_state_dict(torch.load(args.decoder_path)) # If use gpu if torch.cuda.is_available(): encoder.cuda() decoder.cuda() # Prepare Image image_dir = args.image images = os.listdir(image_dir) for image_id in images: if not image_id.endswith('.jpg'): continue image = os.path.join(image_dir, image_id) image = load_image(image, transform) image_tensor = to_var(image, volatile=True) # Generate caption from image try: feature, cnn_features = encoder(image_tensor) sampled_ids = decoder.sample(feature, cnn_features) sampled_ids = sampled_ids.cpu().data.numpy() except: continue # Decode word_ids to words sampled_caption = [] for word_id in sampled_ids: word = vocab.idx2word[word_id] sampled_caption.append(word) if word == '<end>': break sentence = ' '.join(sampled_caption) # Print out image and generated caption. print (image_id + '\t' + sentence)
def main(args): # Image preprocessing transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) # Load vocabulary wrapper with open(args.vocab_path, 'rb') as f: vocab = pickle.load(f) # Build models encoder = EncoderCNN(args.embed_size).eval( ) # eval mode (batchnorm uses moving mean/variance) decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), 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)) images = glob.glob('../data/flickr/train/*.jpg') for i in images: # Prepare an image image = load_image(i, 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 = vocab.idx2word[word_id] sampled_caption.append(word) if word == '<end>': break sentence = ' '.join(sampled_caption) print("{} : {}".format(i[21:], sentence))
def main(cfg): # print(cfg.pretty()) transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(cfg.image.mean, cfg.image.std) ]) print(hydra.utils.to_absolute_path(cfg.train.vocab_path)) with open(hydra.utils.to_absolute_path(cfg.train.vocab_path), 'rb') as f: vocab = pickle.load(f) # モデルの構築 encoder = EncoderCNN(cfg.train.embed_size).eval() decoder = DecoderRNN(cfg.train.embed_size, cfg.train.hidden_size, len(vocab), cfg.train.num_layers) encoder = encoder.to(device) decoder = decoder.to(device) # 学習済みモデルのパラメータを読み込む encoder.load_state_dict( torch.load(hydra.utils.to_absolute_path(cfg.train.encoder_path))) decoder.load_state_dict( torch.load(hydra.utils.to_absolute_path(cfg.train.decoder_path))) # 画像の準備 image = load_image(hydra.utils.to_absolute_path(cfg.sample.image_path), transform) image_tensor = image.to(device) # 入力した画像からキャプションを生成する feature = encoder(image_tensor) # sampled_ids = decoder.sample(feature) sampled_ids = decoder.greedy_decode(features=feature) sampled_ids = sampled_ids[0].cpu().numpy() # word_idsをwordに変換する sampled_caption = [] for word_id in sampled_ids: word = vocab.idx2word[word_id] sampled_caption.append(word) if word == '<end>': break sentence = ' '.join(sampled_caption) print(sentence) image = Image.open(hydra.utils.to_absolute_path(cfg.sample.image_path)) plt.imshow(np.asarray(image))
def main(args): # Image preprocessing transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))]) # Load vocabulary wrapper with open(args.vocab_path, 'rb') as f: vocab = pickle.load(f) # Build models encoder = EncoderCNN(args.embed_size).eval() # eval mode (batchnorm uses moving mean/variance) decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), 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 = vocab.idx2word[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 main(args): # Image preprocessing transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))]) # Load vocabulary wrapper with open(args.vocab_path, 'rb') as f: vocab = pickle.load(f) # Build Models encoder = EncoderCNN(args.embed_size) encoder.eval() # evaluation mode (BN uses moving mean/variance) decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), args.num_layers) # Load the trained model parameters encoder.load_state_dict(torch.load(args.encoder_path)) decoder.load_state_dict(torch.load(args.decoder_path)) # Prepare Image image = load_image(args.image, transform) image_tensor = to_var(image, volatile=True) # If use gpu if torch.cuda.is_available(): encoder.cuda() decoder.cuda() # Generate caption from image feature = encoder(image_tensor) sampled_ids = decoder.sample(feature) sampled_ids = sampled_ids.cpu().data.numpy() # Decode word_ids to words sampled_caption = [] for word_id in sampled_ids: word = vocab.idx2word[word_id] sampled_caption.append(word) if word == '<end>': break sentence = ' '.join(sampled_caption) # Print out image and generated caption. print (sentence) image = Image.open(args.image) plt.imshow(np.asarray(image))
def main(args): # Create model directory if not os.path.exists(args.model_path): os.makedirs(args.model_path) # Image preprocessing, normalization for the pretrained resnet transform = transforms.Compose([ transforms.RandomCrop(args.crop_size), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))]) # Load vocabulary wrapper with open(args.vocab_path, 'rb') as f: vocab = pickle.load(f) # Build data loader data_loader = get_loader(args.image_dir, args.caption_path, vocab, transform, args.batch_size, shuffle=True, num_workers=args.num_workers) # Build the models encoder = EncoderCNN(args.embed_size).to(device) decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), args.num_layers).to(device) # Loss and optimizer criterion = nn.CrossEntropyLoss() params = list(decoder.parameters()) + list(encoder.linear.parameters()) + list(encoder.bn.parameters()) optimizer = torch.optim.Adam(params, lr=args.learning_rate) # Train the models total_step = len(data_loader) for epoch in range(args.num_epochs): for i, (images, captions, lengths) in enumerate(data_loader): # Set mini-batch dataset images = images.to(device) captions = captions.to(device) targets = pack_padded_sequence(captions, lengths, batch_first=True)[0] # Forward, backward and optimize features = encoder(images) outputs = decoder(features, captions, lengths) loss = criterion(outputs, targets) decoder.zero_grad() encoder.zero_grad() loss.backward() optimizer.step() # Print log info if i % args.log_step == 0: print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}, Perplexity: {:5.4f}' .format(epoch, args.num_epochs, i, total_step, loss.item(), np.exp(loss.item()))) # Save the model checkpoints if (i+1) % args.save_step == 0: torch.save(decoder.state_dict(), os.path.join( args.model_path, 'decoder-{}-{}.ckpt'.format(epoch+1, i+1))) torch.save(encoder.state_dict(), os.path.join( args.model_path, 'encoder-{}-{}.ckpt'.format(epoch+1, i+1)))
transforms.Normalize((0.485, 0.456, 0.406), # normalize image for pre-trained model (0.229, 0.224, 0.225))]) # Build data loader. data_loader = get_loader(transform=transform_train, mode='train', batch_size=batch_size, vocab_threshold=vocab_threshold, vocab_from_file=vocab_from_file, cocoapi_loc=COCOPATH) # The size of the vocabulary. vocab_size = len(data_loader.dataset.vocab) # Initialize the encoder and decoder. encoder = EncoderCNN(embed_size) decoder = DecoderRNN(embed_size, hidden_size, vocab_size) # Move models to GPU if CUDA is available. device = torch.device("cuda" if torch.cuda.is_available() else "cpu") encoder.to(device) decoder.to(device) # Define the loss function. criterion = nn.CrossEntropyLoss().cuda() if torch.cuda.is_available() else nn.CrossEntropyLoss() # TODO #3: Specify the learnable parameters of the model. params = list(decoder.parameters()) +\ list(encoder.embed.parameters()) # We don't want to retrain the resnet # TODO #4: Define the optimizer.
(0.229, 0.224, 0.225))]) # Build data loader. data_loader = get_loader(transform=transform_test, mode='test_small', batch_size=batch_size, vocab_threshold=vocab_threshold, vocab_from_file=vocab_from_file, cocoapi_loc=COCOPATH) vocab = data_loader.dataset.vocab # The size of the vocabulary. vocab_size = len(vocab) # Initialize the encoder and decoder. encoder = EncoderCNN(embed_size) decoder = DecoderRNN(embed_size, hidden_size, vocab_size) device = torch.device("cpu") # encoder.to(device) # decoder.to(device) # Load the pretrained model encoder.load_state_dict(torch.load(PRETRAINED_MODEL_PATH.format('encoder'))) decoder.load_state_dict(torch.load(PRETRAINED_MODEL_PATH.format('decoder'))) encoder.eval() decoder.eval() images, conv_images = next(iter(data_loader)) features = encoder(conv_images).unsqueeze(1)
def main(args): # Create model directory if not os.path.exists(args.model_path): os.makedirs(args.model_path) # Image preprocessing # For normalization, see https://github.com/pytorch/vision#models transform = transforms.Compose([ transforms.RandomCrop(args.crop_size), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))]) # Load vocabulary wrapper. with open(args.vocab_path, 'rb') as f: vocab = pickle.load(f) # Build data loader data_loader = get_loader(args.image_dir, args.caption_path, vocab, transform, args.batch_size, shuffle=True, num_workers=args.num_workers) # Build the models encoder = EncoderCNN(args.embed_size) decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), args.num_layers) if torch.cuda.is_available(): encoder.cuda() decoder.cuda() # Loss and Optimizer criterion = nn.CrossEntropyLoss() params = list(decoder.parameters()) + list(encoder.linear.parameters()) + list(encoder.bn.parameters()) optimizer = torch.optim.Adam(params, lr=args.learning_rate) # Train the Models total_step = len(data_loader) for epoch in range(args.num_epochs): for i, (images, captions, lengths) in enumerate(data_loader): # Set mini-batch dataset images = to_var(images, volatile=True) captions = to_var(captions) targets = pack_padded_sequence(captions, lengths, batch_first=True)[0] # Forward, Backward and Optimize decoder.zero_grad() encoder.zero_grad() features = encoder(images) outputs = decoder(features, captions, lengths) loss = criterion(outputs, targets) loss.backward() optimizer.step() # Print log info if i % args.log_step == 0: print('Epoch [%d/%d], Step [%d/%d], Loss: %.4f, Perplexity: %5.4f' %(epoch, args.num_epochs, i, total_step, loss.data[0], np.exp(loss.data[0]))) # Save the models if (i+1) % args.save_step == 0: torch.save(decoder.state_dict(), os.path.join(args.model_path, 'decoder-%d-%d.pkl' %(epoch+1, i+1))) torch.save(encoder.state_dict(), os.path.join(args.model_path, 'encoder-%d-%d.pkl' %(epoch+1, i+1)))