import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from models import *
from location_dataset import LocationDataset
import numpy as np
import math
import utils
parser = argparse.ArgumentParser(description='Location Training with Auxillary Tasks')
parser.add_argument('--data-dir', type=str, default="bedroom", metavar='S')
parser.add_argument('--num-workers', type=int, default=6, metavar='N')
parser.add_argument('--last-epoch', type=int, default=-1, metavar='N') #If positive, use saved epoch
parser.add_argument('--train-size', type=int, default=6400, metavar='N')
parser.add_argument('--save-dir', type=str, default="train/bedroom", metavar='S')
parser.add_argument('--ablation', type=str, default=None, metavar='S')
parser.add_argument('--lr', type=float, default=0.001, metavar='N')
parser.add_argument('--eps', type=float, default=1e-6, metavar='N')
parser.add_argument('--p-auxiliary', type=float, default=0.0, metavar='N')
parser.add_argument('--use-count', action='store_true', default=False) #Use category count if true
parser.add_argument('--no-penalty', action='store_true', default=False) #If True, L_Global is not used
parser.add_argument('--progressive-p', action='store_true', default=False) #If True, start with lower p_auxiliary and gradually increase
args = parser.parse_args()
save_dir = args.save_dir
utils.ensuredir(save_dir)
batch_size = 16
with open(f"data/{args.data_dir}/final_categories_frequency", "r") as f:
lines = f.readlines()
num_categories = len(lines)-2
if args.ablation is None:
num_input_channels = num_categories+9
elif args.ablation == "basic":
num_input_channels = 7
elif args.ablation == "depth":
num_input_channels = 2
else:
raise NotImplementedError
logfile = open(f"{save_dir}/log_location.txt", 'w')
def LOG(msg):
print(msg)
logfile.write(msg + '\n')
logfile.flush()
LOG('Building model...')
model = resnet101(num_classes=num_categories+3, num_input_channels=num_input_channels, use_fc=False)
if args.use_count:
fc = FullyConnected(2048 + num_categories, num_categories+3)
else:
fc = FullyConnected(2048, num_categories+3)
cross_entropy = nn.CrossEntropyLoss()
softmax = nn.Softmax(dim=1)
LOG('Converting to CUDA...')
model.cuda()
fc.cuda()
cross_entropy.cuda()
softmax.cuda()
LOG('Building dataset...')
train_dataset = LocationDataset(
data_root_dir = utils.get_data_root_dir(),
data_dir = args.data_dir,
scene_indices = (0, args.train_size),
p_auxiliary = args.p_auxiliary,
ablation = args.ablation
)
#Size of validation set is 160 by default
validation_dataset = LocationDataset(
data_root_dir = utils.get_data_root_dir(),
data_dir = args.data_dir,
scene_indices = (args.train_size, args.train_size+160),
seed = 42,
p_auxiliary = 0, #Only tests positive examples in validation
ablation = args.ablation
)
LOG('Building data loader...')
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size = batch_size,
num_workers = args.num_workers,
shuffle = True
)
validation_loader = torch.utils.data.DataLoader(
validation_dataset,
batch_size = batch_size,
num_workers = 0,
shuffle = True
)
LOG('Building optimizer...')
optimizer = optim.Adam(list(model.parameters())+list(fc.parameters()),
lr = args.lr,
betas = (0.9,0.999),
eps = args.eps
)
if args.last_epoch < 0:
load = False
starting_epoch = 0
else:
load = True
last_epoch = args.last_epoch
if load:
LOG('Loading saved models...')
model.load_state_dict(torch.load(f"{save_dir}/location_{last_epoch}.pt"))
fc.load_state_dict(torch.load(f"{save_dir}/location_fc_{last_epoch}.pt"))
optimizer.load_state_dict(torch.load(f"{save_dir}/location_optim_backup.pt"))
starting_epoch = last_epoch + 1
current_epoch = starting_epoch
num_seen = 0
if args.progressive_p:
if current_epoch <= 30:
train_dataset.p_auxiliary = 0.0
if current_epoch > 30:
train_dataset.p_auxiliary = 0.5
if current_epoch > 60:
train_dataset.p_auxiliary = 0.7
if current_epoch > 90:
train_dataset.p_auxiliary = 0.9
if current_epoch > 120:
train_dataset.p_auxiliary = 0.95
model.train()
LOG(f'=========================== Epoch {current_epoch} ===========================')
def train():
global num_seen, current_epoch
for batch_idx, (data, t, existing, penalty) \
in enumerate(train_loader):
data, t = data.cuda(), t.cuda()
existing, penalty = existing.cuda(), penalty.cuda()
optimizer.zero_grad()
o_conv = model(data)
if args.use_count:
o_conv = torch.cat([o_conv, existing], 1)
o = fc(o_conv)
loss = cross_entropy(o,t)
if not args.no_penalty:
o_s = softmax(o)[:,0:num_categories]
l_penalty = (o_s * penalty).sum()
if not args.no_penalty:
loss += l_penalty
loss.backward()
optimizer.step()
num_seen += batch_size
if num_seen % 800 == 0:
LOG(f'Examples {num_seen}/10000')
if num_seen % 10000 == 0:
LOG('Validating')
validate()
model.train()
fc.train()
num_seen = 0
current_epoch += 1
LOG(f'=========================== Epoch {current_epoch} ===========================')
LOG(f'{train_dataset.p_auxiliary}')
if current_epoch % 5 == 0:
torch.save(model.state_dict(), f"{save_dir}/location_{current_epoch}.pt")
torch.save(fc.state_dict(), f"{save_dir}/location_fc_{current_epoch}.pt")
torch.save(optimizer.state_dict(), f"{save_dir}/location_optim_backup.pt")
if args.progressive_p:
if current_epoch <= 30:
train_dataset.p_auxiliary = 0.0
if current_epoch > 30:
train_dataset.p_auxiliary = 0.5
if current_epoch > 60:
train_dataset.p_auxiliary = 0.7
if current_epoch > 90:
train_dataset.p_auxiliary = 0.9
if current_epoch > 120:
train_dataset.p_auxiliary = 0.95
def validate():
model.eval()
fc.eval()
total_loss = 0
total_accuracy = 0
for batch_idx, (data, t, existing, penalty) \
in enumerate(validation_loader):
with torch.no_grad():
data, t = data.cuda(), t.cuda()
existing = existing.cuda()
optimizer.zero_grad()
o_conv = model(data)
if args.use_count:
o_conv = torch.cat([o_conv, existing], 1)
o = fc(o_conv)
l = cross_entropy(o,t)
total_loss += l.cpu().data.numpy()
output = softmax(o)
outputs = output.cpu().data.numpy()
targets = t.cpu().data.numpy()
predictions = np.argmax(outputs, axis=1)
num_correct = np.sum(predictions == targets)
total_accuracy += num_correct / batch_size
LOG(f'Loss: {total_loss/10}, Accuracy: {total_accuracy/10}')
while True:
train()
