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) for epoch in range(args.epochs): model.train() log.train(len_dataset=len(dataset.train)) for batch in dataset.train: inputs, targets = (b.to(device) for b in batch) # first forward-backward step
parser.add_argument("--momentum", default=0.9, type=float, help="SGD Momentum.") parser.add_argument("--threads", default=2, type=int, help="Number of CPU threads for dataloaders.") parser.add_argument("--rho", default=2.0, 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, adaptive=args.adaptive, lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) scheduler = StepLR(optimizer, args.learning_rate, args.epochs) for epoch in range(args.epochs): model.train() log.train(len_dataset=len(dataset.train)) for batch in dataset.train: inputs, targets = (b.to(device) for b in batch) # first forward-backward step enable_running_stats(model) predictions = model(inputs) loss = smooth_crossentropy(predictions, targets, smoothing=args.label_smoothing) loss.mean().backward() optimizer.first_step(zero_grad=True)
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) if args.optimizer == "ada_hessian": optimizer = AdaHessian(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay, update_each=args.update_each, average_conv_kernel=args.average_conv_kernel) else: optimizer = SGD(model.parameters(), lr=args.learning_rate, momentum=0.9, nesterov=True, weight_decay=args.weight_decay) scheduler = StepLR(optimizer, args.learning_rate, args.epochs) for epoch in range(args.epochs): model.train()
parser.add_argument("--optimizer", default="ada_hessian", type=str, help="Type of optimizer, supported values are {'ada_hessian', SGD'}.") parser.add_argument("--threads", default=2, type=int, help="Number of CPU threads for dataloaders.") parser.add_argument("--update_each", default=1, type=int, help="Delayed hessian update.") 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) if args.optimizer == "ada_hessian": optimizer = AdaHessian(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay, update_each=args.update_each) else: optimizer = SGD(model.parameters(), lr=args.learning_rate, momentum=0.9, nesterov=True, weight_decay=args.weight_decay) scheduler = StepLR(optimizer, args.learning_rate, args.epochs) for epoch in range(args.epochs): model.train() log.train(len_dataset=len(dataset.train)) for i, (inputs, labels) in enumerate(dataset.train): optimizer.zero_grad() outputs = model(inputs.to(device)) loss = smooth_crossentropy(outputs, labels.to(device)) loss.mean().backward(create_graph=args.optimizer == "ada_hessian") optimizer.step()