# Load model
checkpoint = torch.load(checkpoint)
decoder = checkpoint['decoder']
decoder = decoder.to(device)
decoder.eval()
encoder = checkpoint['encoder']
encoder = encoder.to(device)
encoder.eval()
# Load word map (word2ix)
with open(word_map_file, 'rb') as f:
classlabellist = [line.strip() for line in f]
vocab_size = len(classlabellist)
# customized dataloader
MyDataset = DualLoadDatasets(imgsz, txt_folder, img_folder, bin_folder, split,
Gfiltersz, Gblursigma)
#drop the last batch since it is not divisible by batchsize
test_loader = torch.utils.data.DataLoader(MyDataset,
batch_size=batch_size,
shuffle=False,
num_workers=workers,
pin_memory=True,
drop_last=False)
# Normalization transform
# Applying Transforms to the Data
image_transforms = {
'train':
transforms.Compose([
transforms.ToPILImage(),
transforms.ToTensor(),
def main():
"""
Training and validation.
"""
global checkpoint, start_epoch, fine_tune_encoder
# Initialize / load checkpoint
if checkpoint is None:
encoder = Encoder()
print(encoder)
encoder.fine_tune(fine_tune_encoder)
encoder_optimizer = torch.optim.Adam(
params=filter(lambda p: p.requires_grad, encoder.parameters()),
lr=encoder_lr) if fine_tune_encoder else None
decoder = DecoderWithAttention(attention_dim=attention_dim,
embed_dim=emb_dim,
decoder_dim=decoder_dim,
vocab_size=vocab_size,
encoder_dim=encoder_dim,
dropout=dropout)
decoder_optimizer = torch.optim.Adam(params=filter(
lambda p: p.requires_grad, decoder.parameters()),
lr=decoder_lr)
else:
checkpoint = torch.load(checkpoint)
start_epoch = checkpoint['epoch'] + 1
decoder = checkpoint['decoder']
decoder_optimizer = checkpoint['decoder_optimizer']
encoder = checkpoint['encoder']
encoder_optimizer = checkpoint['encoder_optimizer']
if fine_tune_encoder is True and encoder_optimizer is None:
encoder.fine_tune(fine_tune_encoder)
encoder_optimizer = torch.optim.Adam(params=filter(
lambda p: p.requires_grad, encoder.parameters()),
lr=encoder_lr)
# Move to GPU, if available
decoder = decoder.to(device)
encoder = encoder.to(device)
# Loss function
criterion = nn.CrossEntropyLoss().to(device)
# customized dataloader
MyDataset = DualLoadDatasets(imgsz, txt_folder, img_folder, bin_folder,
split, Gfiltersz, Gblursigma)
#drop the last batch since it is not divisible by batchsize
train_loader = torch.utils.data.DataLoader(MyDataset,
batch_size=batch_size,
shuffle=True,
num_workers=workers,
pin_memory=True,
drop_last=True)
# val_loader = torch.utils.data.DataLoader(
# CaptionDataset(data_folder, data_name, 'VAL', transform=transforms.Compose([normalize])),
# batch_size=batch_size, shuffle=True, num_workers=workers, pin_memory=True)
# Save checkpoint
epoch = 0
save_checkpoint(epoch, encoder, decoder, encoder_optimizer,
decoder_optimizer)
print('saving models to models/checkpoint')
# Epochs
for epoch in range(start_epoch, epochs):
#print(image_transforms)
# One epoch's training
train(train_loader=train_loader,
encoder=encoder,
decoder=decoder,
transform=transform,
criterion=criterion,
encoder_optimizer=encoder_optimizer,
decoder_optimizer=decoder_optimizer,
epoch=epoch)
# Save checkpoint
save_checkpoint(epoch, encoder, decoder, encoder_optimizer,
decoder_optimizer)
print('saving models to models/checkpoint')