def load_dataset(config, augment):
if config.dataset == "mnist":
data_provider = MNIST(config, augment)
elif config.dataset == "imagenette":
data_provider = Imagenette(config, augment)
elif config.dataset == "svhn":
data_provider = Svhn(config, augment)
elif config.dataset == "cifar":
data_provider = Cifar(config, augment)
else:
raise Exception("Unknown dataset %s" % config.dataset)
return data_provider.load_dataset()
type=int,
help="Rho parameter for SAM.")
parser.add_argument("--weight_decay",
default=0.0005,
type=float,
help="L2 weight decay.")
parser.add_argument("--width_factor",
default=8,
type=int,
help="How many times wider compared to normal ResNet.")
args = parser.parse_args()
initialize(args, seed=42)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
dataset = Cifar(args.batch_size, args.threads)
log = Log(log_each=10)
model = WideResNet(args.depth,
args.width_factor,
args.dropout,
in_channels=3,
labels=10).to(device)
base_optimizer = torch.optim.SGD
optimizer = SAM(model.parameters(),
base_optimizer,
rho=args.rho,
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
scheduler = StepLR(optimizer, args.learning_rate, args.epochs)