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()