split='test',
download=True,
transform=transform_test)
else:
train_set = datasets.SVHN(root='../data',
split='train',
download=True,
transform=base_transform)
test_set = datasets.SVHN(root='../data',
split='test',
download=True,
transform=base_transform)
if args.optimizer == 'SGD':
optimizer_factory = optimizer_setup(torch.optim.SGD,
lr=learning_rate,
momentum=0.9)
elif args.optimizer == 'Adam':
optimizer_factory = optimizer_setup(optim.Adam, lr=learning_rate)
if args.model == 'CnnVanilla':
model = CnnVanilla(num_classes=10)
elif args.model == 'AlexNet':
model = AlexNet(num_classes=10)
elif args.model == 'VggNet':
model = VggNet(num_classes=10)
elif args.model == 'ResNet':
model = ResNet(num_classes=10)
elif args.model == 'IFT725Net':
model = IFT725Net(num_classes=10)
elif args.model == 'IFT725UNet':
test_set = HDF5Dataset('test', hdf5_file, transform=acdc_base_transform)
residuts = [
[1,1,1,1,1,1,1,1,1,1,1,1],
[1,0,1,0,1,0,1,0,1,0,1,0],
[1,1,1,1,1,1,0,0,0,0,0,0]
]
num_epochs = 10
best_val_loss = 10
best_residut = [1,1,1,1,1,1,1,1,1,1,1,1]
best_lr = 0
for lr in [0.001, 0.005]:
for residut in residuts:
model = IFT725UNet(num_classes=4, residuts=residut)
optimizer_factory = optimizer_setup(optim.Adam, lr=0.001)
model_trainer = CNNTrainTestManager(model=model,
trainset=train_set,
testset=test_set,
batch_size=20,
loss_fn=nn.CrossEntropyLoss(),
optimizer_factory=optimizer_factory,
validation=0.1,
use_cuda=True)
model_trainer.train(num_epochs)
val_loss = model_trainer.metric_values['val_loss'][-1]
if val_loss < best_val_loss:
best_val_loss = val_loss
best_residut = residut
best_lr = lr
print("best_lr, best_residut:",best_lr, best_residut)