def train():
transform = transforms.Compose([
transforms.Resize((240, 240)),
transforms.RandomCrop(
(224, 224)), #the input size of inception network
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
train_loader, dataset = get_loader(root_folder="archive/Images",
annotation_file="archive/captions.txt",
transform=transform,
batch_size=128,
num_workers=0)
#Set some hyperparamters
torch.backends.cudnn.benchmark = True #Speed up the training process
device = torch.device("cuda" if torch.cuda.is_available() else 'cpu')
load_model = False
save_model = False
train_CNN = False
embed_size = 256
hidden_size = 256
vocab_size = len(dataset.vocab)
num_layers = 1
learning_rate = 3e-4
num_epochs = 100
#for tensorboard
writer = SummaryWriter("runs/flickr")
step = 0
#initialize model, loss etc
model = CNNtoRNN(embed_size, hidden_size, vocab_size,
num_layers).to(device)
# Only finetune the CNN
for name, param in model.EncoderCNN.inception.named_parameters():
if "fc.weight" in name or "fc.bias" in name:
param.requires_grad = True
else:
param.requires_grad = train_CNN
if load_model:
step = load_checkpoint(torch.load("my_checkpoint.pth.tar"), model,
optimizer)
criterion = nn.CrossEntropyLoss(
ignore_index=dataset.vocab.stoi["<PAD>"]) #对于"<PAD>"的词语不需要计算损失
optimizer = optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=learning_rate)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer,
milestones=[60, 120, 140])
model.train()
print('Begins')
imgs, captions = next(iter(train_loader))
for epoch in range(num_epochs):
print_examples(model, device, dataset, save_path='result.txt')
if save_model:
checkpoint = {
"state_dict": model.state_dict(),
"optimizer": optimizer.state_dict(),
"step": step
}
save_checkpoint(checkpoint)
# loop = tqdm(enumerate(train_loader),total=len(train_loader),leave=False)
total_loss = 0
# for idx, (imgs,captions) in loop:
imgs = imgs.to(device)
captions = captions.to(device)
outputs = model(imgs, captions[:-1]) #EOS标志不需要送进网络训练，我们希望他能自己训练出来
# outputs :(seq_len, batch_size, vocabulary_size), 但是交叉熵损失接受二维的tensor
loss = criterion(outputs.reshape(-1, outputs.shape[2]),
captions.reshape(-1))
step += 1
optimizer.zero_grad()
loss.backward(loss)
total_loss += loss.item()
optimizer.step()
print(total_loss)
def train():
transform = transforms.Compose([
transforms.Resize((356, 356)),
transforms.RandomCrop((299, 299)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
train_loader, dataset = get_loader("../data/flickr8k/images/",
"../data/flickr8k/captions.txt",
transform=transform)
torch.backends.cudnn.benchmark = True
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
load_model = False
save_model = True
train_CNN = False
# Hyperparameters
embed_size = 256
hidden_size = 256
vocab_size = len(dataset.vocab)
num_layers = 1
learning_rate = 3e-4
num_epochs = 100
# for tensorboard
writer = SummaryWriter("runs/flickr")
step = 0
# initialize model, loss etc
model = CNNtoRNN(embed_size, hidden_size, vocab_size,
num_layers).to(device)
criterion = nn.CrossEntropyLoss(ignore_index=dataset.vocab.stoi["<PAD>"])
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
# Only finetune the CNN
for name, param in model.encoderCNN.inception.named_parameters():
if "fc.weight" in name or "fc.bias" in name:
param.requires_grad = True
else:
param.requires_grad = train_CNN
if load_model:
step = load_checkpoint(torch.load("my_checkpoint.pth.tar"), model,
optimizer)
model.train()
for epoch in range(num_epochs):
# Uncomment the line below to see a couple of test cases
# print_examples(model, device, dataset)
if save_model:
checkpoint = {
"state_dict": model.state_dict(),
"optimizer": optimizer.state_dict(),
"step": step,
}
save_checkpoint(checkpoint)
for idx, (imgs, captions) in tqdm(enumerate(train_loader),
total=len(train_loader)):
imgs = imgs.to(device)
captions = captions.to(device)
outputs = model(imgs, captions[:-1])
loss = criterion(outputs.reshape(-1, outputs.shape[2]),
captions.reshape(-1))
writer.add_scalar("Training loss", loss.item(), global_step=step)
step += 1
optimizer.zero_grad()
loss.backward(loss)
optimizer.step()
def train():
transform = transforms.Compose([
transforms.Resize((356, 356)),
transforms.RandomCrop((299, 299)), # CNN takes input 299 x 299
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
train_loader, dataset = get_loader(
root_folder='flickr8k/images',
annotation_file='flickr8k/captions.txt',
transform=transform,
num_workers=2,
)
# model configuration
torch.backends.cudnn.benchmark = True
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
load_model = False
save_model = False
train_CNN = False
# Hyperparameters
## We can increase capacity
embed_size = 256
hidden_size = 256
vocab_size = len(dataset.vocab)
num_layers = 1
laerning_rate = 3e-4
num_epochs = 100
# for tensorboard
writer = SummaryWriter('runs/flickr')
step = 0
# initialize model, loss etc
model = CNNtoRNN(embed_size, hidden_size, vocab_size,
num_layers).to(device)
criterion = nn.CrossEntropyLoss(ignore_index=dataset.vocab.stoi["<PAD>"])
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
if load_model:
step = load_checkpoint(
torch.load('my_checkpoint.pth.tar'), model, optimizer
) # we're returning step here so that the loss fucntions continues where it ended
model.train()
for epoch in range(num_epochs):
print_examples(model, device, dataset)
if save_model:
checkpoint = {
"state_dict": model.state_dict(),
"optimizer": optimizer.state_dict(),
"step": step,
}
save_checkpoint(checkpoint)
for idx, (imgs, captions) in enumerate(train_loader):
imgs = imgs.to(device)
captions = captions.to(device)
outputs = model(
imgs, captions[:-1]
) # we actually learn to predict the end token so we're not going to send in the end token
loss = criterion(
outputs.reshape(-1, outputs.shape[2]), captions.reshape(-1)
) #predicting for each example we're predicting for a bunch of different time steps
# example , 20 words that it's predicting and then each word has its logit corresponding to each word in the vocabulary right here.
## so we have three dimensions here , but the criterion only 2 dimension
### output -> (seq_len, N, vocabulary_size) , target -> (seq_len , N)
writer.add_scalar("Training loss", loss.item(), global_step=step)
step += 1
optimizer.zero_grad()
loss.backward(loss)
optimizer.step()
from get_loader import get_loader
import torchvision.transforms as transforms
import pickle
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
])
train_loader, dataset = get_loader(
root_folder="flickr8k/images",
annotation_file="flickr8k/captions.txt",
transform=transform,
num_workers=2,
)
with open('vocab_itos.pkl', 'wb') as f:
pickle.dump(dataset.vocab.itos, f)
def train():
transform = transforms.Compose([
transforms.Resize((356, 356)),
transforms.RandomCrop((299, 299)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
train_loader, dataset = get_loader(
root_folder=
"/mnt/liguanlin/DataSets/ImageCaptionDatasets/flickr8k/images",
annotation_file=
"/mnt/liguanlin/DataSets/ImageCaptionDatasets/flickr8k/captions.txt",
transform=transform,
num_workers=2,
)
torch.backends.cudnn.benchmark = True
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
load_model = False
save_model = True
#Hyperparameters
embed_size = 256
hidden_size = 256
vocab_size = len(dataset.vocab)
num_layers = 1
learning_rate = 3e-4
num_epochs = 100
#for tensorboard
writer = SummaryWriter("runs/flickr")
step = 0
#initialize model, loss etc
model = CNNtoRNN(embed_size, hidden_size, vocab_size,
num_layers).to(device)
criterion = nn.CrossEntropyLoss(ignore_index=dataset.vocab.stoi["<PAD>"])
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
if load_model:
step = load_checkpoint(torch.load("my_checkpoint.pth.tar"), model,
optimizer)
model.train()
for epoch in range(num_epochs):
print_examples(model, device, dataset)
if save_model:
checkpoint = {
"state_dict": model.state_dict(),
"optimizer": optimizer.state_dict(),
"step": step,
}
save_checkpoint(checkpoint)
for idx, (imgs, captions) in enumerate(train_loader):
imgs = imgs.to(device)
captions = captions.to(device)
outputs = model(imgs, captions[:-1])
loss = criterion(outputs.reshape(-1, outputs.shape[2]),
captions.reshape(-1))
#record loss
writer.add_scalar("Training loss", loss.item(), global_step=step)
step += 1
optimizer.zero_grad()
loss.backward()
optimizer.step()
def inferrence(model, dataset, image):
transform = transforms.Compose([
transforms.Resize((299, 299)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
model.eval()
image = transform(image).unsqueeze(0).to(
'cuda' if torch.cuda.is_available() else 'cpu')
image_predict = model.caption_image(image, dataset.vocab)
print("Predicted :" + " ".join(image_predict))
if __name__ == "__main__":
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
test_img = Image.open("test_examples/footable.jpg").convert("RGB")
_, dataset = get_loader(root_folder="archive/Images",
annotation_file="archive/captions.txt",
transform=None,
batch_size=64,
num_workers=0)
embed_size = 256
hidden_size = 256
vocab_size = len(dataset.vocab)
num_layers = 1
model = CNNtoRNN(embed_size, hidden_size, vocab_size,
num_layers).to(device)
model.load_state_dict(torch.load("my_checkpoint.pth.tar")["state_dict"])
model.eval()
inferrence(model, dataset, test_img)
np.random.seed(seed)
cuda = torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
print(device)
N, N_te = min(args.dataset_size, 60000), min(args.dataset_size, 10000)
args.dataset_size = N
args.test_size = N_te
experiment.log_multiple_params(vars(args))
experiment.disable_mp()
args.device = device
now = datetime.datetime.now()
train_loader, test_loader = get_loader(args)
if args.model == 'cnn':
model = SmallCNN().to(device)
elif args.model == 'cnn_bn':
model = SmallCNN_BN().to(device)
elif args.model == 'logreg':
model = LogReg().to(device)
elif args.model == 'mlp':
model = MLP(args.hidden_size, args.activation).to(device)
elif args.model == 'big_mlp':
model = BigMLP(args.hidden_size).to(device)
else:
print('No model recognized')
optimizer = torch.optim.SGD(model.parameters(),