def main(args, conf):
## DIY dataset
tr_dataset = AudioDataset(conf, data)
cv_dataset = AudioDataset(conf, data)
## DIY dataloader
batch_size = conf.get("batch_size", 33)
nun_workers = conf.get("num_workers", 5)
tr_loader = AudioDataLoader(tr_dataset,
batch_size=batch_size,
num_workers=nun_workers)
cv_loader = AudioDataLoader(cv_dataset,
batch_size=batch_size,
num_workers=nun_workers)
loader_dict = {'tr_loader': tr_loader, 'cv_loader': cv_loader}
## DIY model
model = SegNet(conf)
print(model)
model.cuda()
optimizier = get_optim(model, conf)
# solver
solver = Solver(conf, loader_dict, model, optimizier)
solver.train()
def main():
global args, best_prec1
args = parser.parse_args()
# Check if the save directory exists or not
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
model = SegNet(3, 3)
#model.features = torch.nn.DataParallel(model.features)
if use_gpu:
model.cuda()
# 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'])
#optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.evaluate, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# data_transforms = {
# 'train': transforms.Compose([
# transforms.Scale(256),
# transforms.RandomSizedCrop(224),
# transforms.RandomHorizontalFlip(),
# transforms.ToTensor(),
# transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]),
# ]),
# 'val': transforms.Compose([
# transforms.Scale(256),
# transforms.CenterCrop(224),
# transforms.RandomHorizontalFlip(),
# transforms.ToTensor(),
# transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]),
# ]),
# }
data_transforms = {
'train':
transforms.Compose([
transforms.Scale((224, 224)),
transforms.ToTensor(),
]),
'val':
transforms.Compose([
transforms.Scale((224, 224)),
transforms.ToTensor(),
]),
}
data_dir = '/media/salman/DATA/NUST/MS RIME/Thesis/MICCAI Dataset/miccai_all_images'
image_datasets = {
x: miccaiDataset(os.path.join(data_dir, x), data_transforms[x])
for x in ['train', 'val']
}
dataloaders = {
x: torch.utils.data.DataLoader(image_datasets[x],
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers)
for x in ['train', 'val']
}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
# Define loss function (criterion) and optimizer
criterion = nn.MSELoss().cuda()
if args.half:
model.half()
criterion.half()
#optimizer = torch.optim.SGD(model.parameters(), args.lr,
# momentum=args.momentum,
# weight_decay=args.weight_decay)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
if args.evaluate:
validate(dataloaders['val'], model, criterion)
return
for epoch in range(args.start_epoch, args.epochs):
#adjust_learning_rate(optimizer, epoch)
# Train for one epoch
train(dataloaders['train'], model, criterion, optimizer, epoch)
# Evaulate on validation set
prec1 = validate(dataloaders['val'], model, criterion)
prec1 = prec1.cpu().data.numpy()
# Remember best prec1 and save checkpoint
print(prec1)
print(best_prec1)
is_best = prec1 < best_prec1
best_prec1 = min(prec1, best_prec1)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
#'optimizer': optimizer.state_dict(),
},
is_best,
filename=os.path.join(args.save_dir,
'checkpoint_{}.tar'.format(epoch)))
