chore: import upstream snapshot with attribution
@@ -0,0 +1,148 @@
|
||||
import os
|
||||
import torch
|
||||
import torchvision
|
||||
import torch.nn as nn
|
||||
from torchvision import transforms
|
||||
from torchvision.utils import save_image
|
||||
|
||||
|
||||
# Device configuration
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
# Hyper-parameters
|
||||
latent_size = 64
|
||||
hidden_size = 256
|
||||
image_size = 784
|
||||
num_epochs = 200
|
||||
batch_size = 100
|
||||
sample_dir = 'samples'
|
||||
|
||||
# Create a directory if not exists
|
||||
if not os.path.exists(sample_dir):
|
||||
os.makedirs(sample_dir)
|
||||
|
||||
# Image processing
|
||||
# transform = transforms.Compose([
|
||||
# transforms.ToTensor(),
|
||||
# transforms.Normalize(mean=(0.5, 0.5, 0.5), # 3 for RGB channels
|
||||
# std=(0.5, 0.5, 0.5))])
|
||||
transform = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.5], # 1 for greyscale channels
|
||||
std=[0.5])])
|
||||
|
||||
# MNIST dataset
|
||||
mnist = torchvision.datasets.MNIST(root='../../data/',
|
||||
train=True,
|
||||
transform=transform,
|
||||
download=True)
|
||||
|
||||
# Data loader
|
||||
data_loader = torch.utils.data.DataLoader(dataset=mnist,
|
||||
batch_size=batch_size,
|
||||
shuffle=True)
|
||||
|
||||
# Discriminator
|
||||
D = nn.Sequential(
|
||||
nn.Linear(image_size, hidden_size),
|
||||
nn.LeakyReLU(0.2),
|
||||
nn.Linear(hidden_size, hidden_size),
|
||||
nn.LeakyReLU(0.2),
|
||||
nn.Linear(hidden_size, 1),
|
||||
nn.Sigmoid())
|
||||
|
||||
# Generator
|
||||
G = nn.Sequential(
|
||||
nn.Linear(latent_size, hidden_size),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_size, hidden_size),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_size, image_size),
|
||||
nn.Tanh())
|
||||
|
||||
# Device setting
|
||||
D = D.to(device)
|
||||
G = G.to(device)
|
||||
|
||||
# Binary cross entropy loss and optimizer
|
||||
criterion = nn.BCELoss()
|
||||
d_optimizer = torch.optim.Adam(D.parameters(), lr=0.0002)
|
||||
g_optimizer = torch.optim.Adam(G.parameters(), lr=0.0002)
|
||||
|
||||
def denorm(x):
|
||||
out = (x + 1) / 2
|
||||
return out.clamp(0, 1)
|
||||
|
||||
def reset_grad():
|
||||
d_optimizer.zero_grad()
|
||||
g_optimizer.zero_grad()
|
||||
|
||||
# Start training
|
||||
total_step = len(data_loader)
|
||||
for epoch in range(num_epochs):
|
||||
for i, (images, _) in enumerate(data_loader):
|
||||
images = images.reshape(batch_size, -1).to(device)
|
||||
|
||||
# Create the labels which are later used as input for the BCE loss
|
||||
real_labels = torch.ones(batch_size, 1).to(device)
|
||||
fake_labels = torch.zeros(batch_size, 1).to(device)
|
||||
|
||||
# ================================================================== #
|
||||
# Train the discriminator #
|
||||
# ================================================================== #
|
||||
|
||||
# Compute BCE_Loss using real images where BCE_Loss(x, y): - y * log(D(x)) - (1-y) * log(1 - D(x))
|
||||
# Second term of the loss is always zero since real_labels == 1
|
||||
outputs = D(images)
|
||||
d_loss_real = criterion(outputs, real_labels)
|
||||
real_score = outputs
|
||||
|
||||
# Compute BCELoss using fake images
|
||||
# First term of the loss is always zero since fake_labels == 0
|
||||
z = torch.randn(batch_size, latent_size).to(device)
|
||||
fake_images = G(z)
|
||||
outputs = D(fake_images)
|
||||
d_loss_fake = criterion(outputs, fake_labels)
|
||||
fake_score = outputs
|
||||
|
||||
# Backprop and optimize
|
||||
d_loss = d_loss_real + d_loss_fake
|
||||
reset_grad()
|
||||
d_loss.backward()
|
||||
d_optimizer.step()
|
||||
|
||||
# ================================================================== #
|
||||
# Train the generator #
|
||||
# ================================================================== #
|
||||
|
||||
# Compute loss with fake images
|
||||
z = torch.randn(batch_size, latent_size).to(device)
|
||||
fake_images = G(z)
|
||||
outputs = D(fake_images)
|
||||
|
||||
# We train G to maximize log(D(G(z)) instead of minimizing log(1-D(G(z)))
|
||||
# For the reason, see the last paragraph of section 3. https://arxiv.org/pdf/1406.2661.pdf
|
||||
g_loss = criterion(outputs, real_labels)
|
||||
|
||||
# Backprop and optimize
|
||||
reset_grad()
|
||||
g_loss.backward()
|
||||
g_optimizer.step()
|
||||
|
||||
if (i+1) % 200 == 0:
|
||||
print('Epoch [{}/{}], Step [{}/{}], d_loss: {:.4f}, g_loss: {:.4f}, D(x): {:.2f}, D(G(z)): {:.2f}'
|
||||
.format(epoch, num_epochs, i+1, total_step, d_loss.item(), g_loss.item(),
|
||||
real_score.mean().item(), fake_score.mean().item()))
|
||||
|
||||
# Save real images
|
||||
if (epoch+1) == 1:
|
||||
images = images.reshape(images.size(0), 1, 28, 28)
|
||||
save_image(denorm(images), os.path.join(sample_dir, 'real_images.png'))
|
||||
|
||||
# Save sampled images
|
||||
fake_images = fake_images.reshape(fake_images.size(0), 1, 28, 28)
|
||||
save_image(denorm(fake_images), os.path.join(sample_dir, 'fake_images-{}.png'.format(epoch+1)))
|
||||
|
||||
# Save the model checkpoints
|
||||
torch.save(G.state_dict(), 'G.ckpt')
|
||||
torch.save(D.state_dict(), 'D.ckpt')
|
||||
@@ -0,0 +1,59 @@
|
||||
# Image Captioning
|
||||
The goal of image captioning is to convert a given input image into a natural language description. The encoder-decoder framework is widely used for this task. The image encoder is a convolutional neural network (CNN). In this tutorial, we used [resnet-152](https://arxiv.org/abs/1512.03385) model pretrained on the [ILSVRC-2012-CLS](http://www.image-net.org/challenges/LSVRC/2012/) image classification dataset. The decoder is a long short-term memory (LSTM) network.
|
||||
|
||||

|
||||
|
||||
#### Training phase
|
||||
For the encoder part, the pretrained CNN extracts the feature vector from a given input image. The feature vector is linearly transformed to have the same dimension as the input dimension of the LSTM network. For the decoder part, source and target texts are predefined. For example, if the image description is **"Giraffes standing next to each other"**, the source sequence is a list containing **['\<start\>', 'Giraffes', 'standing', 'next', 'to', 'each', 'other']** and the target sequence is a list containing **['Giraffes', 'standing', 'next', 'to', 'each', 'other', '\<end\>']**. Using these source and target sequences and the feature vector, the LSTM decoder is trained as a language model conditioned on the feature vector.
|
||||
|
||||
#### Test phase
|
||||
In the test phase, the encoder part is almost same as the training phase. The only difference is that batchnorm layer uses moving average and variance instead of mini-batch statistics. This can be easily implemented using [encoder.eval()](https://github.com/yunjey/pytorch-tutorial/blob/master/tutorials/03-advanced/image_captioning/sample.py#L37). For the decoder part, there is a significant difference between the training phase and the test phase. In the test phase, the LSTM decoder can't see the image description. To deal with this problem, the LSTM decoder feeds back the previosly generated word to the next input. This can be implemented using a [for-loop](https://github.com/yunjey/pytorch-tutorial/blob/master/tutorials/03-advanced/image_captioning/model.py#L48).
|
||||
|
||||
|
||||
|
||||
## Usage
|
||||
|
||||
|
||||
#### 1. Clone the repositories
|
||||
```bash
|
||||
git clone https://github.com/pdollar/coco.git
|
||||
cd coco/PythonAPI/
|
||||
make
|
||||
python setup.py build
|
||||
python setup.py install
|
||||
cd ../../
|
||||
git clone https://github.com/yunjey/pytorch-tutorial.git
|
||||
cd pytorch-tutorial/tutorials/03-advanced/image_captioning/
|
||||
```
|
||||
|
||||
#### 2. Download the dataset
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
chmod +x download.sh
|
||||
./download.sh
|
||||
```
|
||||
|
||||
#### 3. Preprocessing
|
||||
|
||||
```bash
|
||||
python build_vocab.py
|
||||
python resize.py
|
||||
```
|
||||
|
||||
#### 4. Train the model
|
||||
|
||||
```bash
|
||||
python train.py
|
||||
```
|
||||
|
||||
#### 5. Test the model
|
||||
|
||||
```bash
|
||||
python sample.py --image='png/example.png'
|
||||
```
|
||||
|
||||
<br>
|
||||
|
||||
## Pretrained model
|
||||
If you do not want to train the model from scratch, you can use a pretrained model. You can download the pretrained model [here](https://www.dropbox.com/s/ne0ixz5d58ccbbz/pretrained_model.zip?dl=0) and the vocabulary file [here](https://www.dropbox.com/s/26adb7y9m98uisa/vocap.zip?dl=0). You should extract pretrained_model.zip to `./models/` and vocab.pkl to `./data/` using `unzip` command.
|
||||
@@ -0,0 +1,76 @@
|
||||
import nltk
|
||||
import pickle
|
||||
import argparse
|
||||
from collections import Counter
|
||||
from pycocotools.coco import COCO
|
||||
|
||||
|
||||
class Vocabulary(object):
|
||||
"""Simple vocabulary wrapper."""
|
||||
def __init__(self):
|
||||
self.word2idx = {}
|
||||
self.idx2word = {}
|
||||
self.idx = 0
|
||||
|
||||
def add_word(self, word):
|
||||
if not word in self.word2idx:
|
||||
self.word2idx[word] = self.idx
|
||||
self.idx2word[self.idx] = word
|
||||
self.idx += 1
|
||||
|
||||
def __call__(self, word):
|
||||
if not word in self.word2idx:
|
||||
return self.word2idx['<unk>']
|
||||
return self.word2idx[word]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.word2idx)
|
||||
|
||||
def build_vocab(json, threshold):
|
||||
"""Build a simple vocabulary wrapper."""
|
||||
coco = COCO(json)
|
||||
counter = Counter()
|
||||
ids = coco.anns.keys()
|
||||
for i, id in enumerate(ids):
|
||||
caption = str(coco.anns[id]['caption'])
|
||||
tokens = nltk.tokenize.word_tokenize(caption.lower())
|
||||
counter.update(tokens)
|
||||
|
||||
if (i+1) % 1000 == 0:
|
||||
print("[{}/{}] Tokenized the captions.".format(i+1, len(ids)))
|
||||
|
||||
# If the word frequency is less than 'threshold', then the word is discarded.
|
||||
words = [word for word, cnt in counter.items() if cnt >= threshold]
|
||||
|
||||
# Create a vocab wrapper and add some special tokens.
|
||||
vocab = Vocabulary()
|
||||
vocab.add_word('<pad>')
|
||||
vocab.add_word('<start>')
|
||||
vocab.add_word('<end>')
|
||||
vocab.add_word('<unk>')
|
||||
|
||||
# Add the words to the vocabulary.
|
||||
for i, word in enumerate(words):
|
||||
vocab.add_word(word)
|
||||
return vocab
|
||||
|
||||
def main(args):
|
||||
vocab = build_vocab(json=args.caption_path, threshold=args.threshold)
|
||||
vocab_path = args.vocab_path
|
||||
with open(vocab_path, 'wb') as f:
|
||||
pickle.dump(vocab, f)
|
||||
print("Total vocabulary size: {}".format(len(vocab)))
|
||||
print("Saved the vocabulary wrapper to '{}'".format(vocab_path))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--caption_path', type=str,
|
||||
default='data/annotations/captions_train2014.json',
|
||||
help='path for train annotation file')
|
||||
parser.add_argument('--vocab_path', type=str, default='./data/vocab.pkl',
|
||||
help='path for saving vocabulary wrapper')
|
||||
parser.add_argument('--threshold', type=int, default=4,
|
||||
help='minimum word count threshold')
|
||||
args = parser.parse_args()
|
||||
main(args)
|
||||
@@ -0,0 +1,105 @@
|
||||
import torch
|
||||
import torchvision.transforms as transforms
|
||||
import torch.utils.data as data
|
||||
import os
|
||||
import pickle
|
||||
import numpy as np
|
||||
import nltk
|
||||
from PIL import Image
|
||||
from build_vocab import Vocabulary
|
||||
from pycocotools.coco import COCO
|
||||
|
||||
|
||||
class CocoDataset(data.Dataset):
|
||||
"""COCO Custom Dataset compatible with torch.utils.data.DataLoader."""
|
||||
def __init__(self, root, json, vocab, transform=None):
|
||||
"""Set the path for images, captions and vocabulary wrapper.
|
||||
|
||||
Args:
|
||||
root: image directory.
|
||||
json: coco annotation file path.
|
||||
vocab: vocabulary wrapper.
|
||||
transform: image transformer.
|
||||
"""
|
||||
self.root = root
|
||||
self.coco = COCO(json)
|
||||
self.ids = list(self.coco.anns.keys())
|
||||
self.vocab = vocab
|
||||
self.transform = transform
|
||||
|
||||
def __getitem__(self, index):
|
||||
"""Returns one data pair (image and caption)."""
|
||||
coco = self.coco
|
||||
vocab = self.vocab
|
||||
ann_id = self.ids[index]
|
||||
caption = coco.anns[ann_id]['caption']
|
||||
img_id = coco.anns[ann_id]['image_id']
|
||||
path = coco.loadImgs(img_id)[0]['file_name']
|
||||
|
||||
image = Image.open(os.path.join(self.root, path)).convert('RGB')
|
||||
if self.transform is not None:
|
||||
image = self.transform(image)
|
||||
|
||||
# Convert caption (string) to word ids.
|
||||
tokens = nltk.tokenize.word_tokenize(str(caption).lower())
|
||||
caption = []
|
||||
caption.append(vocab('<start>'))
|
||||
caption.extend([vocab(token) for token in tokens])
|
||||
caption.append(vocab('<end>'))
|
||||
target = torch.Tensor(caption)
|
||||
return image, target
|
||||
|
||||
def __len__(self):
|
||||
return len(self.ids)
|
||||
|
||||
|
||||
def collate_fn(data):
|
||||
"""Creates mini-batch tensors from the list of tuples (image, caption).
|
||||
|
||||
We should build custom collate_fn rather than using default collate_fn,
|
||||
because merging caption (including padding) is not supported in default.
|
||||
|
||||
Args:
|
||||
data: list of tuple (image, caption).
|
||||
- image: torch tensor of shape (3, 256, 256).
|
||||
- caption: torch tensor of shape (?); variable length.
|
||||
|
||||
Returns:
|
||||
images: torch tensor of shape (batch_size, 3, 256, 256).
|
||||
targets: torch tensor of shape (batch_size, padded_length).
|
||||
lengths: list; valid length for each padded caption.
|
||||
"""
|
||||
# Sort a data list by caption length (descending order).
|
||||
data.sort(key=lambda x: len(x[1]), reverse=True)
|
||||
images, captions = zip(*data)
|
||||
|
||||
# Merge images (from tuple of 3D tensor to 4D tensor).
|
||||
images = torch.stack(images, 0)
|
||||
|
||||
# Merge captions (from tuple of 1D tensor to 2D tensor).
|
||||
lengths = [len(cap) for cap in captions]
|
||||
targets = torch.zeros(len(captions), max(lengths)).long()
|
||||
for i, cap in enumerate(captions):
|
||||
end = lengths[i]
|
||||
targets[i, :end] = cap[:end]
|
||||
return images, targets, lengths
|
||||
|
||||
def get_loader(root, json, vocab, transform, batch_size, shuffle, num_workers):
|
||||
"""Returns torch.utils.data.DataLoader for custom coco dataset."""
|
||||
# COCO caption dataset
|
||||
coco = CocoDataset(root=root,
|
||||
json=json,
|
||||
vocab=vocab,
|
||||
transform=transform)
|
||||
|
||||
# Data loader for COCO dataset
|
||||
# This will return (images, captions, lengths) for each iteration.
|
||||
# images: a tensor of shape (batch_size, 3, 224, 224).
|
||||
# captions: a tensor of shape (batch_size, padded_length).
|
||||
# lengths: a list indicating valid length for each caption. length is (batch_size).
|
||||
data_loader = torch.utils.data.DataLoader(dataset=coco,
|
||||
batch_size=batch_size,
|
||||
shuffle=shuffle,
|
||||
num_workers=num_workers,
|
||||
collate_fn=collate_fn)
|
||||
return data_loader
|
||||
@@ -0,0 +1,11 @@
|
||||
mkdir data
|
||||
wget http://msvocds.blob.core.windows.net/annotations-1-0-3/captions_train-val2014.zip -P ./data/
|
||||
wget http://images.cocodataset.org/zips/train2014.zip -P ./data/
|
||||
wget http://images.cocodataset.org/zips/val2014.zip -P ./data/
|
||||
|
||||
unzip ./data/captions_train-val2014.zip -d ./data/
|
||||
rm ./data/captions_train-val2014.zip
|
||||
unzip ./data/train2014.zip -d ./data/
|
||||
rm ./data/train2014.zip
|
||||
unzip ./data/val2014.zip -d ./data/
|
||||
rm ./data/val2014.zip
|
||||
@@ -0,0 +1,56 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision.models as models
|
||||
from torch.nn.utils.rnn import pack_padded_sequence
|
||||
|
||||
|
||||
class EncoderCNN(nn.Module):
|
||||
def __init__(self, embed_size):
|
||||
"""Load the pretrained ResNet-152 and replace top fc layer."""
|
||||
super(EncoderCNN, self).__init__()
|
||||
resnet = models.resnet152(pretrained=True)
|
||||
modules = list(resnet.children())[:-1] # delete the last fc layer.
|
||||
self.resnet = nn.Sequential(*modules)
|
||||
self.linear = nn.Linear(resnet.fc.in_features, embed_size)
|
||||
self.bn = nn.BatchNorm1d(embed_size, momentum=0.01)
|
||||
|
||||
def forward(self, images):
|
||||
"""Extract feature vectors from input images."""
|
||||
with torch.no_grad():
|
||||
features = self.resnet(images)
|
||||
features = features.reshape(features.size(0), -1)
|
||||
features = self.bn(self.linear(features))
|
||||
return features
|
||||
|
||||
|
||||
class DecoderRNN(nn.Module):
|
||||
def __init__(self, embed_size, hidden_size, vocab_size, num_layers, max_seq_length=20):
|
||||
"""Set the hyper-parameters and build the layers."""
|
||||
super(DecoderRNN, self).__init__()
|
||||
self.embed = nn.Embedding(vocab_size, embed_size)
|
||||
self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True)
|
||||
self.linear = nn.Linear(hidden_size, vocab_size)
|
||||
self.max_seg_length = max_seq_length
|
||||
|
||||
def forward(self, features, captions, lengths):
|
||||
"""Decode image feature vectors and generates captions."""
|
||||
embeddings = self.embed(captions)
|
||||
embeddings = torch.cat((features.unsqueeze(1), embeddings), 1)
|
||||
packed = pack_padded_sequence(embeddings, lengths, batch_first=True)
|
||||
hiddens, _ = self.lstm(packed)
|
||||
outputs = self.linear(hiddens[0])
|
||||
return outputs
|
||||
|
||||
def sample(self, features, states=None):
|
||||
"""Generate captions for given image features using greedy search."""
|
||||
sampled_ids = []
|
||||
inputs = features.unsqueeze(1)
|
||||
for i in range(self.max_seg_length):
|
||||
hiddens, states = self.lstm(inputs, states) # hiddens: (batch_size, 1, hidden_size)
|
||||
outputs = self.linear(hiddens.squeeze(1)) # outputs: (batch_size, vocab_size)
|
||||
_, predicted = outputs.max(1) # predicted: (batch_size)
|
||||
sampled_ids.append(predicted)
|
||||
inputs = self.embed(predicted) # inputs: (batch_size, embed_size)
|
||||
inputs = inputs.unsqueeze(1) # inputs: (batch_size, 1, embed_size)
|
||||
sampled_ids = torch.stack(sampled_ids, 1) # sampled_ids: (batch_size, max_seq_length)
|
||||
return sampled_ids
|
||||
|
After Width: | Height: | Size: 220 KiB |
|
After Width: | Height: | Size: 246 KiB |
|
After Width: | Height: | Size: 246 KiB |
@@ -0,0 +1,5 @@
|
||||
matplotlib
|
||||
nltk
|
||||
numpy
|
||||
Pillow
|
||||
argparse
|
||||
@@ -0,0 +1,42 @@
|
||||
import argparse
|
||||
import os
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def resize_image(image, size):
|
||||
"""Resize an image to the given size."""
|
||||
return image.resize(size, Image.ANTIALIAS)
|
||||
|
||||
def resize_images(image_dir, output_dir, size):
|
||||
"""Resize the images in 'image_dir' and save into 'output_dir'."""
|
||||
if not os.path.exists(output_dir):
|
||||
os.makedirs(output_dir)
|
||||
|
||||
images = os.listdir(image_dir)
|
||||
num_images = len(images)
|
||||
for i, image in enumerate(images):
|
||||
with open(os.path.join(image_dir, image), 'r+b') as f:
|
||||
with Image.open(f) as img:
|
||||
img = resize_image(img, size)
|
||||
img.save(os.path.join(output_dir, image), img.format)
|
||||
if (i+1) % 100 == 0:
|
||||
print ("[{}/{}] Resized the images and saved into '{}'."
|
||||
.format(i+1, num_images, output_dir))
|
||||
|
||||
def main(args):
|
||||
image_dir = args.image_dir
|
||||
output_dir = args.output_dir
|
||||
image_size = [args.image_size, args.image_size]
|
||||
resize_images(image_dir, output_dir, image_size)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--image_dir', type=str, default='./data/train2014/',
|
||||
help='directory for train images')
|
||||
parser.add_argument('--output_dir', type=str, default='./data/resized2014/',
|
||||
help='directory for saving resized images')
|
||||
parser.add_argument('--image_size', type=int, default=256,
|
||||
help='size for image after processing')
|
||||
args = parser.parse_args()
|
||||
main(args)
|
||||
@@ -0,0 +1,81 @@
|
||||
import torch
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import argparse
|
||||
import pickle
|
||||
import os
|
||||
from torchvision import transforms
|
||||
from build_vocab import Vocabulary
|
||||
from model import EncoderCNN, DecoderRNN
|
||||
from PIL import Image
|
||||
|
||||
|
||||
# Device configuration
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
def load_image(image_path, transform=None):
|
||||
image = Image.open(image_path).convert('RGB')
|
||||
image = image.resize([224, 224], Image.LANCZOS)
|
||||
|
||||
if transform is not None:
|
||||
image = transform(image).unsqueeze(0)
|
||||
|
||||
return 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))
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--image', type=str, required=True, help='input image for generating caption')
|
||||
parser.add_argument('--encoder_path', type=str, default='models/encoder-5-3000.pkl', help='path for trained encoder')
|
||||
parser.add_argument('--decoder_path', type=str, default='models/decoder-5-3000.pkl', help='path for trained decoder')
|
||||
parser.add_argument('--vocab_path', type=str, default='data/vocab.pkl', help='path for vocabulary wrapper')
|
||||
|
||||
# Model parameters (should be same as paramters in train.py)
|
||||
parser.add_argument('--embed_size', type=int , default=256, help='dimension of word embedding vectors')
|
||||
parser.add_argument('--hidden_size', type=int , default=512, help='dimension of lstm hidden states')
|
||||
parser.add_argument('--num_layers', type=int , default=1, help='number of layers in lstm')
|
||||
args = parser.parse_args()
|
||||
main(args)
|
||||
@@ -0,0 +1,101 @@
|
||||
import argparse
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import numpy as np
|
||||
import os
|
||||
import pickle
|
||||
from data_loader import get_loader
|
||||
from build_vocab import Vocabulary
|
||||
from model import EncoderCNN, DecoderRNN
|
||||
from torch.nn.utils.rnn import pack_padded_sequence
|
||||
from torchvision import transforms
|
||||
|
||||
|
||||
# Device configuration
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
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)))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--model_path', type=str, default='models/' , help='path for saving trained models')
|
||||
parser.add_argument('--crop_size', type=int, default=224 , help='size for randomly cropping images')
|
||||
parser.add_argument('--vocab_path', type=str, default='data/vocab.pkl', help='path for vocabulary wrapper')
|
||||
parser.add_argument('--image_dir', type=str, default='data/resized2014', help='directory for resized images')
|
||||
parser.add_argument('--caption_path', type=str, default='data/annotations/captions_train2014.json', help='path for train annotation json file')
|
||||
parser.add_argument('--log_step', type=int , default=10, help='step size for prining log info')
|
||||
parser.add_argument('--save_step', type=int , default=1000, help='step size for saving trained models')
|
||||
|
||||
# Model parameters
|
||||
parser.add_argument('--embed_size', type=int , default=256, help='dimension of word embedding vectors')
|
||||
parser.add_argument('--hidden_size', type=int , default=512, help='dimension of lstm hidden states')
|
||||
parser.add_argument('--num_layers', type=int , default=1, help='number of layers in lstm')
|
||||
|
||||
parser.add_argument('--num_epochs', type=int, default=5)
|
||||
parser.add_argument('--batch_size', type=int, default=128)
|
||||
parser.add_argument('--num_workers', type=int, default=2)
|
||||
parser.add_argument('--learning_rate', type=float, default=0.001)
|
||||
args = parser.parse_args()
|
||||
print(args)
|
||||
main(args)
|
||||
@@ -0,0 +1,33 @@
|
||||
# Neural Style Transfer
|
||||
|
||||
[Neural style transfer](https://arxiv.org/abs/1508.06576) is an algorithm that combines the content of one image with the style of another image using CNN. Given a content image and a style image, the goal is to generate a target image that minimizes the content difference with the content image and the style difference with the style image.
|
||||
|
||||
<p align="center"><img width="100%" src="png/neural_style2.png" /></p>
|
||||
|
||||
|
||||
#### Content loss
|
||||
|
||||
To minimize the content difference, we forward propagate the content image and the target image to pretrained [VGGNet](https://arxiv.org/abs/1409.1556) respectively, and extract feature maps from multiple convolutional layers. Then, the target image is updated to minimize the [mean-squared error](https://github.com/yunjey/pytorch-tutorial/blob/master/tutorials/03-advanced/neural_style_transfer/main.py#L81-L82) between the feature maps of the content image and its feature maps.
|
||||
|
||||
#### Style loss
|
||||
|
||||
As in computing the content loss, we forward propagate the style image and the target image to the VGGNet and extract convolutional feature maps. To generate a texture that matches the style of the style image, we update the target image by minimizing the mean-squared error between the Gram matrix of the style image and the Gram matrix of the target image (feature correlation minimization). See [here](https://github.com/yunjey/pytorch-tutorial/blob/master/tutorials/03-advanced/neural_style_transfer/main.py#L84-L94) for how to compute the style loss.
|
||||
|
||||
|
||||
|
||||
|
||||
<br>
|
||||
|
||||
## Usage
|
||||
|
||||
```bash
|
||||
$ pip install -r requirements.txt
|
||||
$ python main.py --content='png/content.png' --style='png/style.png'
|
||||
```
|
||||
|
||||
<br>
|
||||
|
||||
## Results
|
||||
The following is the result of applying variaous styles of artwork to Anne Hathaway's photograph.
|
||||
|
||||

|
||||
@@ -0,0 +1,126 @@
|
||||
from __future__ import division
|
||||
from torchvision import models
|
||||
from torchvision import transforms
|
||||
from PIL import Image
|
||||
import argparse
|
||||
import torch
|
||||
import torchvision
|
||||
import torch.nn as nn
|
||||
import numpy as np
|
||||
|
||||
|
||||
# Device configuration
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
def load_image(image_path, transform=None, max_size=None, shape=None):
|
||||
"""Load an image and convert it to a torch tensor."""
|
||||
image = Image.open(image_path)
|
||||
|
||||
if max_size:
|
||||
scale = max_size / max(image.size)
|
||||
size = np.array(image.size) * scale
|
||||
image = image.resize(size.astype(int), Image.ANTIALIAS)
|
||||
|
||||
if shape:
|
||||
image = image.resize(shape, Image.LANCZOS)
|
||||
|
||||
if transform:
|
||||
image = transform(image).unsqueeze(0)
|
||||
|
||||
return image.to(device)
|
||||
|
||||
|
||||
class VGGNet(nn.Module):
|
||||
def __init__(self):
|
||||
"""Select conv1_1 ~ conv5_1 activation maps."""
|
||||
super(VGGNet, self).__init__()
|
||||
self.select = ['0', '5', '10', '19', '28']
|
||||
self.vgg = models.vgg19(pretrained=True).features
|
||||
|
||||
def forward(self, x):
|
||||
"""Extract multiple convolutional feature maps."""
|
||||
features = []
|
||||
for name, layer in self.vgg._modules.items():
|
||||
x = layer(x)
|
||||
if name in self.select:
|
||||
features.append(x)
|
||||
return features
|
||||
|
||||
|
||||
def main(config):
|
||||
|
||||
# Image preprocessing
|
||||
# VGGNet was trained on ImageNet where images are normalized by mean=[0.485, 0.456, 0.406] and std=[0.229, 0.224, 0.225].
|
||||
# We use the same normalization statistics here.
|
||||
transform = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=(0.485, 0.456, 0.406),
|
||||
std=(0.229, 0.224, 0.225))])
|
||||
|
||||
# Load content and style images
|
||||
# Make the style image same size as the content image
|
||||
content = load_image(config.content, transform, max_size=config.max_size)
|
||||
style = load_image(config.style, transform, shape=[content.size(2), content.size(3)])
|
||||
|
||||
# Initialize a target image with the content image
|
||||
target = content.clone().requires_grad_(True)
|
||||
|
||||
optimizer = torch.optim.Adam([target], lr=config.lr, betas=[0.5, 0.999])
|
||||
vgg = VGGNet().to(device).eval()
|
||||
|
||||
for step in range(config.total_step):
|
||||
|
||||
# Extract multiple(5) conv feature vectors
|
||||
target_features = vgg(target)
|
||||
content_features = vgg(content)
|
||||
style_features = vgg(style)
|
||||
|
||||
style_loss = 0
|
||||
content_loss = 0
|
||||
for f1, f2, f3 in zip(target_features, content_features, style_features):
|
||||
# Compute content loss with target and content images
|
||||
content_loss += torch.mean((f1 - f2)**2)
|
||||
|
||||
# Reshape convolutional feature maps
|
||||
_, c, h, w = f1.size()
|
||||
f1 = f1.view(c, h * w)
|
||||
f3 = f3.view(c, h * w)
|
||||
|
||||
# Compute gram matrix
|
||||
f1 = torch.mm(f1, f1.t())
|
||||
f3 = torch.mm(f3, f3.t())
|
||||
|
||||
# Compute style loss with target and style images
|
||||
style_loss += torch.mean((f1 - f3)**2) / (c * h * w)
|
||||
|
||||
# Compute total loss, backprop and optimize
|
||||
loss = content_loss + config.style_weight * style_loss
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
if (step+1) % config.log_step == 0:
|
||||
print ('Step [{}/{}], Content Loss: {:.4f}, Style Loss: {:.4f}'
|
||||
.format(step+1, config.total_step, content_loss.item(), style_loss.item()))
|
||||
|
||||
if (step+1) % config.sample_step == 0:
|
||||
# Save the generated image
|
||||
denorm = transforms.Normalize((-2.12, -2.04, -1.80), (4.37, 4.46, 4.44))
|
||||
img = target.clone().squeeze()
|
||||
img = denorm(img).clamp_(0, 1)
|
||||
torchvision.utils.save_image(img, 'output-{}.png'.format(step+1))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--content', type=str, default='png/content.png')
|
||||
parser.add_argument('--style', type=str, default='png/style.png')
|
||||
parser.add_argument('--max_size', type=int, default=400)
|
||||
parser.add_argument('--total_step', type=int, default=2000)
|
||||
parser.add_argument('--log_step', type=int, default=10)
|
||||
parser.add_argument('--sample_step', type=int, default=500)
|
||||
parser.add_argument('--style_weight', type=float, default=100)
|
||||
parser.add_argument('--lr', type=float, default=0.003)
|
||||
config = parser.parse_args()
|
||||
print(config)
|
||||
main(config)
|
||||
|
After Width: | Height: | Size: 599 KiB |
|
After Width: | Height: | Size: 1.3 MiB |
|
After Width: | Height: | Size: 493 KiB |
|
After Width: | Height: | Size: 682 KiB |
|
After Width: | Height: | Size: 947 KiB |
|
After Width: | Height: | Size: 1.2 MiB |
|
After Width: | Height: | Size: 1.9 MiB |
@@ -0,0 +1,4 @@
|
||||
argparse
|
||||
torch
|
||||
torchvision
|
||||
Pillow
|
||||
@@ -0,0 +1,101 @@
|
||||
import os
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torchvision
|
||||
from torchvision import transforms
|
||||
from torchvision.utils import save_image
|
||||
|
||||
|
||||
# Device configuration
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
# Create a directory if not exists
|
||||
sample_dir = 'samples'
|
||||
if not os.path.exists(sample_dir):
|
||||
os.makedirs(sample_dir)
|
||||
|
||||
# Hyper-parameters
|
||||
image_size = 784
|
||||
h_dim = 400
|
||||
z_dim = 20
|
||||
num_epochs = 15
|
||||
batch_size = 128
|
||||
learning_rate = 1e-3
|
||||
|
||||
# MNIST dataset
|
||||
dataset = torchvision.datasets.MNIST(root='../../data',
|
||||
train=True,
|
||||
transform=transforms.ToTensor(),
|
||||
download=True)
|
||||
|
||||
# Data loader
|
||||
data_loader = torch.utils.data.DataLoader(dataset=dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=True)
|
||||
|
||||
|
||||
# VAE model
|
||||
class VAE(nn.Module):
|
||||
def __init__(self, image_size=784, h_dim=400, z_dim=20):
|
||||
super(VAE, self).__init__()
|
||||
self.fc1 = nn.Linear(image_size, h_dim)
|
||||
self.fc2 = nn.Linear(h_dim, z_dim)
|
||||
self.fc3 = nn.Linear(h_dim, z_dim)
|
||||
self.fc4 = nn.Linear(z_dim, h_dim)
|
||||
self.fc5 = nn.Linear(h_dim, image_size)
|
||||
|
||||
def encode(self, x):
|
||||
h = F.relu(self.fc1(x))
|
||||
return self.fc2(h), self.fc3(h)
|
||||
|
||||
def reparameterize(self, mu, log_var):
|
||||
std = torch.exp(log_var/2)
|
||||
eps = torch.randn_like(std)
|
||||
return mu + eps * std
|
||||
|
||||
def decode(self, z):
|
||||
h = F.relu(self.fc4(z))
|
||||
return F.sigmoid(self.fc5(h))
|
||||
|
||||
def forward(self, x):
|
||||
mu, log_var = self.encode(x)
|
||||
z = self.reparameterize(mu, log_var)
|
||||
x_reconst = self.decode(z)
|
||||
return x_reconst, mu, log_var
|
||||
|
||||
model = VAE().to(device)
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
|
||||
|
||||
# Start training
|
||||
for epoch in range(num_epochs):
|
||||
for i, (x, _) in enumerate(data_loader):
|
||||
# Forward pass
|
||||
x = x.to(device).view(-1, image_size)
|
||||
x_reconst, mu, log_var = model(x)
|
||||
|
||||
# Compute reconstruction loss and kl divergence
|
||||
# For KL divergence, see Appendix B in VAE paper or http://yunjey47.tistory.com/43
|
||||
reconst_loss = F.binary_cross_entropy(x_reconst, x, size_average=False)
|
||||
kl_div = - 0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
|
||||
|
||||
# Backprop and optimize
|
||||
loss = reconst_loss + kl_div
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
if (i+1) % 10 == 0:
|
||||
print ("Epoch[{}/{}], Step [{}/{}], Reconst Loss: {:.4f}, KL Div: {:.4f}"
|
||||
.format(epoch+1, num_epochs, i+1, len(data_loader), reconst_loss.item(), kl_div.item()))
|
||||
|
||||
with torch.no_grad():
|
||||
# Save the sampled images
|
||||
z = torch.randn(batch_size, z_dim).to(device)
|
||||
out = model.decode(z).view(-1, 1, 28, 28)
|
||||
save_image(out, os.path.join(sample_dir, 'sampled-{}.png'.format(epoch+1)))
|
||||
|
||||
# Save the reconstructed images
|
||||
out, _, _ = model(x)
|
||||
x_concat = torch.cat([x.view(-1, 1, 28, 28), out.view(-1, 1, 28, 28)], dim=3)
|
||||
save_image(x_concat, os.path.join(sample_dir, 'reconst-{}.png'.format(epoch+1)))
|
||||