def main():
# load trained model
print("*** loading model from {model}".format(model=args.model))
model = modified_resnet50()
if args.cuda:
model = model.cuda()
checkpoint = torch.load(args.model)
model.load_state_dict(checkpoint['state_dict'])
# calculate output
df = eval_one_dir(args.img_dir, model)
# write csv file
df.to_csv(args.output_csvpath, index=False)
def main():
# load trained model
print("*** loading model from {model}".format(model=args.model))
model = modified_resnet50()
if args.cuda:
model = model.cuda()
'''with open(args.model) as f:
model.load_state_dict( torch.load(f)['state_dict'])
print("*** calculating the model output of the images in {img_dir}"
.format(img_dir = args.img_dir))
'''
checkpoint = torch.load(args.model)
model.load_state_dict(checkpoint['state_dict'])
#optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
#epoch = checkpoint['epoch']
#loss = checkpoint['loss']error in create table BRANCH,
# calculate output
df = eval_one_dir(args.img_dir, model)
# write csv file
df.to_csv(args.output_csvpath, index=False)
def main():
global best_loss
loss_history_train = []
loss_history_val = []
data_dir = args.data_dir
img_dir_train = (r"train")
img_dir_val = (r"test")
txt_file_train = (r"annot_train.txt")
txt_file_val = (r"annot_test.txt")
# load pretrained resnet50 with a modified last fully connected layer
model = modified_resnet50()
# we need three different criterion for training
criterion_protest = nn.BCELoss()
criterion_violence = nn.MSELoss()
criterion_visattr = nn.BCELoss()
criterions = [criterion_protest, criterion_violence, criterion_visattr]
if args.cuda and not torch.cuda.is_available():
raise Exception("No GPU Found")
if args.cuda:
model = model.cuda()
criterions = [criterion.cuda() for criterion in criterions]
# we are not training the frozen layers
parameters = filter(lambda p: p.requires_grad, model.parameters())
optimizer = torch.optim.SGD(
parameters, args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay
)
#Incase your program terminates and you want to start from the last checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_loss = checkpoint['best_loss']
args.start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
loss_history_train = checkpoint['loss_history_train']
loss_history_val = checkpoint['loss_history_val']
if args.change_lr:
for param_group in optimizer.param_groups:
param_group['lr'] = args.lr
else:
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
eigval = torch.Tensor([0.2175, 0.0188, 0.0045])
eigvec = torch.Tensor([[-0.5675, 0.7192, 0.4009],
[-0.5808, -0.0045, -0.8140],
[-0.5836, -0.6948, 0.4203]])
train_dataset = ProtestDataset(
txt_file = txt_file_train,
img_dir = img_dir_train,
transform = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomRotation(30),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(
brightness = 0.4,
contrast = 0.4,
saturation = 0.4,
),
transforms.ToTensor(),
Lighting(0.1, eigval, eigvec),
normalize,
]))
val_dataset = ProtestDataset(
txt_file = txt_file_val,
img_dir = img_dir_val,
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
train_loader = DataLoader(
train_dataset,
num_workers = args.workers,
batch_size = args.batch_size,
shuffle = True
)
val_loader = DataLoader(
val_dataset,
num_workers = args.workers,
batch_size = args.batch_size)
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
loss_history_train_this = train(train_loader, model, criterions, optimizer, epoch)
loss_val, loss_history_val_this = validate(val_loader, model,
criterions, epoch)
loss_history_train.append(loss_history_train_this)
loss_history_val.append(loss_history_val_this)
# loss = loss_val.avg
is_best = loss_val < best_loss
if is_best:
print('best model!!')
best_loss = min(loss_val, best_loss)
save_checkpoint({
'epoch' : epoch + 1,
'state_dict' : model.state_dict(),
'best_loss' : best_loss,
'optimizer' : optimizer.state_dict(),
'loss_history_train': loss_history_train,
'loss_history_val': loss_history_val
}, is_best)