def main():
global args, best_prec1, USE_GPU, device
args = parser.parse_args()
with open(args.config) as f:
config = yaml.load(f)
for k, v in config['common'].items():
setattr(args, k, v)
## Set random seeds ##
torch.manual_seed(args.random_seed)
np.random.seed(args.random_seed)
random.seed(args.random_seed)
#print(args)
#print(xxx)
# create models
if args.input_size != 224 or args.image_size != 256:
image_size = args.image_size
input_size = args.input_size
else:
image_size = 256
input_size = 224
print("Input image size: {}, test size: {}".format(image_size, input_size))
if "model" in config.keys():
model = models.__dict__[args.arch](**config['model'])
else:
model = models.__dict__[args.arch]()
device = torch.device(
'cuda:' + str(args.gpus[0]) if torch.cuda.is_available() else "cpu")
str_input_size = '1x3x224x224'
if args.summary:
input_size = tuple(int(x) for x in str_input_size.split('x'))
stat(model, input_size)
return
if USE_GPU:
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.random_seed)
args.gpus = [int(i) for i in args.gpus.split(',')]
model = torch.nn.DataParallel(model, device_ids=args.gpus)
model.to(device)
count_params(model)
pytorch_total_params = sum(p.numel() for p in model.parameters())
print('total_params', pytorch_total_params)
# define loss function (criterion) and optimizer
criterion = FocalLoss(device, 2, gamma=args.fl_gamma)
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
if args.speed:
input_size = tuple(int(x) for x in str_input_size.split('x'))
iteration = 1000
compute_speed(model, input_size, device, iteration)
return
##########################################################
# adjust start learning rate
args.lr = args.lr * (args.batch_size // 32)
#########################################################
# optionally resume from a checkpoint
if args.resume:
print(os.getcwd())
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'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
# Data loading code
normalize = transforms.Normalize(
mean=[0.14300402, 0.1434545,
0.14277956], ##accorcoding to casia-surf val to commpute
std=[0.10050353, 0.100842826, 0.10034215])
img_size = args.input_size
ratio = 224.0 / float(img_size)
train_dataset = CASIA(
args,
transforms.Compose([
transforms.RandomResizedCrop(img_size),
transforms.RandomHorizontalFlip(),
#ColorTransform(p=0.3),
transforms.ToTensor(),
ColorAugmentation(),
normalize,
]),
phase_train=True)
val_dataset = CASIA(args,
transforms.Compose([
transforms.Resize(int(256 * ratio)),
transforms.CenterCrop(img_size),
transforms.ToTensor(),
normalize,
]),
phase_train=False,
phase_test=args.phase_test)
train_sampler = None
val_sampler = None
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=(train_sampler is None),
sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True,
sampler=val_sampler)
if args.evaluate:
validate(val_loader, model, criterion, args.start_epoch)
return
else:
print(model)
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(val_loader, model, criterion, epoch)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
if is_best:
print('epoch: {} The best is {} last best is {}'.format(
epoch, prec1, best_prec1))
best_prec1 = max(prec1, best_prec1)
if not os.path.exists(args.save_path):
os.makedirs(args.save_path)
if is_best:
save_name = '{}/{}_{}_best.pth.tar'.format(args.save_path,
args.model_name, epoch)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
},
filename=save_name)
def main():
global args, best_prec1, USE_GPU
args = parser.parse_args()
with open(args.config) as f:
config = yaml.load(f)
for k, v in config['common'].items():
setattr(args, k, v)
# create models
if args.input_size != 224 or args.image_size != 256:
image_size = args.image_size
input_size = args.input_size
else:
image_size = 256
input_size = 224
print("Input image size: {}, test size: {}".format(image_size, input_size))
if "model" in config.keys():
model = models.__dict__[args.arch](**config['model'])
else:
model = models.__dict__[args.arch]()
if USE_GPU:
model = model.cuda()
model = torch.nn.DataParallel(model)
count_params(model)
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# 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)
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# Data loading code
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
img_size = args.input_size
ratio = 64.0 / float(img_size)
train_dataset = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(img_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
ColorAugmentation(),
normalize,
]))
val_dataset = datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Resize(int(256 * ratio)),
transforms.CenterCrop(img_size),
transforms.ToTensor(),
normalize,
]))
# if args.distributed:
# train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
# val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset)
# else:
train_sampler = None
val_sampler = None
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=(train_sampler is None),
sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True,
sampler=val_sampler)
if args.evaluate:
validate(val_loader, model, criterion)
return
for epoch in range(args.start_epoch, args.epochs):
# if args.distributed:
# train_sampler.set_epoch(epoch)
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
if not os.path.exists(args.save_path):
os.mkdir(args.save_path)
save_name = '{}/{}_{}_best.pth.tar'.format(args.save_path, args.model_name, epoch) if is_best else\
'{}/{}_{}.pth.tar'.format(args.save_path, args.model_name, epoch)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
},
filename=save_name)
def main():
global args, best_prec1, USE_GPU, run
args = parser.parse_args()
with open(args.config) as f:
config = yaml.load(f)
for k, v in config["common"].items():
setattr(args, k, v)
# create models
if args.input_size != 224 or args.image_size != 256:
image_size = args.image_size
input_size = args.input_size
else:
image_size = 256
input_size = 224
print("Input image size: {}, test size: {}".format(image_size, input_size))
if "model" in config.keys():
model = models.__dict__[args.arch](**config["model"])
else:
model = models.__dict__[args.arch]()
if USE_GPU:
model = model.cuda()
model = torch.nn.DataParallel(model)
run = wandb.init(project="fishnet", config=args)
count_params(model)
# pms.summary(
# model,
# torch.zeros(1, 3, 224, 224),
# max_depth=3,
# show_parent_layers=True,
# show_hierarchical=True,
# print_summary=True,
# )
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# 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)
best_prec1 = checkpoint["best_prec1"]
model.load_state_dict(checkpoint["state_dict"])
optimizer.load_state_dict(checkpoint["optimizer"])
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# Data loading code
traindir = os.path.join(args.data, "train")
valdir = os.path.join(args.data, "val")
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
img_size = args.input_size
ratio = 224.0 / float(img_size)
train_dataset = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(img_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
ColorAugmentation(),
normalize,
]),
)
val_dataset = datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Resize(int(256 * ratio)),
transforms.CenterCrop(img_size),
transforms.ToTensor(),
normalize,
]),
)
# print(train_dataset[0])
# if args.distributed:
# train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
# val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset)
# else:
train_sampler = None
val_sampler = None
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=False, # (train_sampler is None),
sampler=train_sampler,
)
val_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=args.batch_size // 2,
shuffle=False,
num_workers=args.workers,
pin_memory=False, # True,
sampler=val_sampler,
)
if args.evaluate:
validate(val_loader, model, criterion)
return
for epoch in range(args.start_epoch, args.epochs):
# if args.distributed:
# train_sampler.set_epoch(epoch)
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
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 not os.path.exists(args.save_path):
os.mkdir(args.save_path)
save_name = ("{}/{}_{}_best.pth.tar".format(args.save_path,
args.model_name, epoch)
if is_best else "{}/{}_{}.pth.tar".format(
args.save_path, args.model_name, epoch))
save_checkpoint(
{
"epoch": epoch + 1,
"arch": args.arch,
"state_dict": model.state_dict(),
"best_prec1": best_prec1,
"optimizer": optimizer.state_dict(),
},
filename=save_name,
)
print(f"epoch {epoch}: Prec@1 {prec1} BestPrec@1 {best_prec1}")
run.finish()