def main():
print(args)
print("=> creating model '{}'".format(args.arch))
model = Ensemble()
model = torch.nn.DataParallel(model).cuda()
print(model)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
test_data = datautil.SceneDataset(args.data,img_transform=
transforms.Compose([
transforms.Resize((args.img_size,args.img_size)),
transforms.ToTensor(),
normalize]))
test_loader = torch.utils.data.DataLoader(test_data,batch_size=args.batch_size,shuffle=False,num_workers=4,pin_memory=True)
checkpoint = torch.load(args.test_model)
model.load_state_dict(checkpoint['state_dict'])
#model.load_state_dict(checkpoint)
if os.path.isdir(args.data):
ret = test(test_loader,model)
imgs = [i[:-4] for i in os.listdir(args.data)]
with open('result3_.csv', 'w') as f:
'''
f.write(','.join(['FILE_ID','CATEGORY_ID'])+'\n')
f.write('\n'.join([','.join([str(a),str(b)]) for a,b in zip(imgs,ret)]))
'''
#FILE_ID,CATEGORY_ID0,CATEGORY_ID1,CATEGORY_ID2
f.write(','.join(['FILE_ID','CATEGORY_ID0','CATEGORY_ID1','CATEGORY_ID2'])+'\n')
f.write('\n'.join([','.join([str(a)]+[str(int(i)) for i in b]) for a,b in zip(imgs,ret)]))
else:
test_labeled(test_loader,model)
def main():
global best_prec1
if args.tensorboard:
configure('log/'+args.arch.lower() + '_bs' + str(args.batch_size) + '_ep' + str(args.epochs) + '_loglr' + str(args.lr) +
'_size' + str(args.img_size)+ '_wd' + str(args.weight_decay))
print(args)
# create model
print("=> creating model '{}'".format(args.arch))
#if args.arch.lower().startswith('resnet'):
# model.avgpool = nn.AvgPool2d(args.img_size // 32, 1)
#model.fc = nn.Linear(model.fc.in_features, args.num_classes)
# default parameter n_class=1000, input_size=224, width_mult=1.
model = Ensemble()
if not args.resume:
model.MobileNetV2.load_state_dict(torch.load('mobilenet_pretrained.pth'))
model.NASNetAMobile.load_state_dict(torch.load('nasnet_pretrained.pth'))
model = torch.nn.DataParallel(model).cuda()
print(model)
# optionally resume from a 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_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# Data loading code
import datautil
'''
norm_dict = {
#320:transforms.Normalize(mean=[0.4333,0.4429,0.4313],std=[ 1., 1., 1.]),
320:transforms.Normalize(mean=[0.4333,0.4429,0.4313],std=[0.2295, 0.2385, 0.2479]),
0:transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225]),
256:transforms.Normalize(mean=[0.4333,0.4429,0.4313],std=[ 0.2295, 0.2385, 0.2479]),
224:transforms.Normalize(mean=[0.4333,0.4429,0.4313],std=[ 0.2295, 0.2385, 0.2479]),
}
norm_default = norm_dict[0]
normalize = norm_dict[args.img_size]
currrent
tensor([[ 0.4828, 0.4693, 0.4602]], device='cuda:0')
tensor([[ 45.3332, 41.1241, 45.7719]], device='cuda:0')
'''
#normalize = transforms.Normalize(mean=[0.48280172,0.46929353,0.46019437],std=[0.25859008,0.28414325,0.288328])
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_data = datautil.SceneDataset(args.data,img_transform=
transforms.Compose([
transforms.RandomResizedCrop(args.img_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize]))
train_loader = torch.utils.data.DataLoader(train_data,batch_size=args.batch_size,shuffle=True,num_workers=args.workers,pin_memory=True)
if args.val:
val_data = datautil.SceneDataset(args.val, img_transform=
transforms.Compose([
#transforms.Scale(256),
transforms.Resize((args.img_size,args.img_size)),
transforms.ToTensor(),
normalize]))
val_loader = torch.utils.data.DataLoader(val_data, batch_size=args.batch_size//2, shuffle=False,
num_workers=args.workers, pin_memory=True)
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
# optimizer = torch.optim.Adam(model.parameters(), args.lr, weight_decay=args.weight_decay)
optimizer = torch.optim.SGD(model.parameters(), args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
#optimizer = torch.optim.RMSprop(model.parameters(), args.lr,
# momentum=args.momentum,
# weight_decay=args.weight_decay,eps=1)
if args.evaluate:
validate(val_loader, model, criterion,0)
return
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
if args.val:
prec1 = validate(val_loader, model, criterion,epoch)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
if epoch % args.interval == 0:
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
})
