global args, best_prec1, model, train_dataset, test_loader

args = parser.parse_args()

args.distributed = args.world_size > 1

if args.distributed:

dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,

world_size=args.world_size)

logging.basicConfig(

format="%(message)s",

handlers=[

logging.FileHandler("{0}/{1}.log".format(args.log, sys.argv[0].replace('.py','') + datetime.now().strftime('_%H_%M_%d_%m_%Y'))),

logging.StreamHandler()

],

level=logging.INFO)

logging.info('creating model and allocating memory')

model = SocialNet(fashion=args.fashion, event=args.event, num_class=config['num_class']).cuda()

logging.info('initilzing model with pretrained weights')

model._initialize_weights(args.event, args.fashion)

# define loss function (criterion) and optimizer

criterion = nn.CrossEntropyLoss().cuda()

# optimizer = torch.optim.SGD(model.parameters(), args.lr,

# momentum=args.momentum,

# weight_decay=args.weight_decay)

param = list(model.event.parameters()) + list(model.fc.parameters())

optimizer = torch.optim.Adam(param, args.lr,

betas=(0.9, 0.999), eps=1e-08, weight_decay=args.weight_decay)

# optionally resume from a checkpoint

cudnn.benchmark = True

# Data loading code

logging.info('creating data loaders')

train_dataset = Dataset(config, mode='train')

test_dataset = Dataset(config, mode='eval')

if args.resume:

if os.path.isfile(args.resume):

logging.info("=> loading checkpoint '{}'".format(args.resume))

checkpoint = torch.load(args.resume)

args.start_epoch = checkpoint['epoch']

best_prec1 = checkpoint['best_prec1']

model.load_state_dict(checkpoint['state_dict'])

optimizer.load_state_dict(checkpoint['optimizer'])

logging.info("=> loaded checkpoint '{}' (epoch {})"

.format(args.resume, checkpoint['epoch']))

else:

logging.info("=> no checkpoint found at '{}'".format(args.resume))

train_loader = torch.utils.data.DataLoader(

train_dataset, batch_size=args.batch_size, shuffle=True,

num_workers=args.workers, pin_memory=True)

test_loader = torch.utils.data.DataLoader(

test_dataset, batch_size=args.batch_size, shuffle=True,

num_workers=args.workers, pin_memory=True)

if args.evaluate:

validate(val_loader, model, criterion)

return

logging.info('start training')

for epoch in range(args.start_epoch, args.epochs):

adjust_learning_rate(optimizer, epoch)

# train for one epoch

train(train_loader, model, criterion, optimizer, epoch)

# evaluate on validation set

prec1 = validate(test_loader, model, criterion)

# remember best prec@1 and save checkpoint

is_best = prec1 > best_prec1

best_prec1 = max(prec1, best_prec1)

save_checkpoint({

'epoch': epoch + 1,

'state_dict': model.state_dict(),

'best_prec1': best_prec1,

'optimizer' : optimizer.state_dict(),

batch_time = AverageMeter()

data_time = AverageMeter()

losses = AverageMeter()

top1 = AverageMeter()

top5 = AverageMeter()

# switch to train mode

model.train()

end = time.time()

global img_var, img_1_var, img_2_var, target_var

for i, (img, img_1, img_2, target) in enumerate(train_loader):

# measure data loading time

data_time.update(time.time() - end)

target = target.cuda(async=True)

img_var = torch.autograd.Variable(img).cuda()

img_1_var = torch.autograd.Variable(img_1).cuda()

img_2_var = torch.autograd.Variable(img_2).cuda()

target_var = torch.autograd.Variable(target)

# compute output

output = model(img_var, img_1_var, img_2_var)

loss = criterion(output, target_var)

# measure accuracy and record loss

prec1, prec5, correct, count = accuracy(output.data, target, \

topk=(1, 5), num_cls=train_loader.dataset.config['num_class'])

losses.update(loss.data[0], img.size(0))

top1.update(prec1[0], img.size(0))

top5.update(prec5[0], img.size(0))

# compute gradient and do SGD step

optimizer.zero_grad()

loss.backward()

optimizer.step()

# measure elapsed time

batch_time.update(time.time() - end)

end = time.time()

if i % args.print_freq == 0:

logging.info('Epoch: [{0}][{1}/{2}]\t'

'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'

'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'

'Loss {loss.val:.4f} ({loss.avg:.4f})\t'

'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'

'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(

epoch, i, len(train_loader), batch_time=batch_time,

batch_time = AverageMeter()

losses = AverageMeter()

top1 = AverageMeter()

top5 = AverageMeter()

# switch to evaluate mode

model.eval()

end = time.time()

correct = list()

count = list()

for i, (img, img_1, img_2, target) in enumerate(val_loader):

# measure data loading time

target = target.cuda(async=True)

img_var = torch.autograd.Variable(img, volatile=True).cuda()

img_1_var = torch.autograd.Variable(img_1, volatile=True).cuda()

img_2_var = torch.autograd.Variable(img_2, volatile=True).cuda()

target_var = torch.autograd.Variable(target, volatile=True)

# compute output

output = model(img_var, img_1_var, img_2_var)

loss = criterion(output, target_var)

# measure accuracy and record loss

prec1, prec5, correct_batch, count_batch = accuracy(output.data, target, \

topk=(1, 5), num_cls=val_loader.dataset.config['num_class'])

correct.append(correct_batch)

count.append(count_batch)

losses.update(loss.data[0], img.size(0))

top1.update(prec1[0], img.size(0))

top5.update(prec5[0], img.size(0))

# measure elapsed time

batch_time.update(time.time() - end)

end = time.time()

if i % args.print_freq == 0:

logging.info('Test: [{0}/{1}]\t'

'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'

'Loss {loss.val:.4f} ({loss.avg:.4f})\t'

'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'

'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(

i, len(val_loader), batch_time=batch_time, loss=losses,

top1=top1, top5=top5))

logging.info(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'

.format(top1=top1, top5=top5))

correct = np.array(correct).sum(0)

count = np.array(count).sum(0).astype(np.int32)

pres = correct/count

for idx , item in enumerate(pres):

msg = str(idx) + ' ' + str(count[idx]) + ' ' + str(item)

logging.info(msg)

filename = os.path.join(args.weights, filename)

torch.save(state, filename)

if is_best:

logging.warning('***********************saving best model *********************')

best = os.path.join(args.weights, '{0}-{1}-best.pth.tar'.format(sys.argv[0][:-3], config['num_class']))

"""Computes and stores the average and current value"""

def __init__(self):

self.reset()

def reset(self):

self.val = 0

self.avg = 0

self.sum = 0

self.count = 0

def update(self, val, n=1):

self.val = val

self.sum += val * n

self.count += n

"""Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""

lr = args.lr * (0.1 ** (epoch // 30))

for param_group in optimizer.param_groups:

"""Computes the precision@k for the specified values of k"""

maxk = max(topk)

batch_size = target.size(0)

_, pred = output.topk(maxk, 1, True, True)

pred = pred.t()

correct = pred.eq(target.view(1, -1).expand_as(pred))

res = []

for k in topk:

correct_k = correct[:k]

res.append(correct_k.view(-1).float().sum(0, keepdim=True).mul_(100.0 / batch_size))

cls_correct = np.zeros(num_cls)

cls_count = np.zeros(num_cls)

for idx, tar in enumerate(target):

if correct[0,idx]:

cls_correct[target[idx]] += 1

cls_count[target[idx]] += 1

res.append(cls_correct)

res.append(cls_count)