# check to see if stuff has run already
if not args.ignore_previous:
with open(args.save_file, 'r') as f:
for line in f.readlines():
if args.meta_data == line.split(',')[0]:
print(f"Metadata {args.meta_data} already ran. Exiting.")
exit()
# training setup
train_loader, test_loader_eval, train_loader_eval, num_classes = get_data(
args)
if args.model == 'fc':
if args.dataset == 'mnist':
net = fc(width=args.width,
depth=args.depth,
num_classes=num_classes).to(args.device)
elif args.dataset == 'cifar10':
net = fc(width=args.width,
depth=args.depth,
num_classes=num_classes,
input_dim=3 * 32 * 32).to(args.device)
elif args.model == 'alexnet':
if args.dataset == 'mnist':
net = alexnet(input_height=28,
input_width=28,
input_channels=1,
num_classes=num_classes)
else:
net = alexnet(ch=args.scale,
num_classes=num_classes).to(args.device)
if torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
print(args)
# training setup
train_loader, test_loader_eval, train_loader_eval, input_dim, num_classes = get_data(
args)
if args.model == 'fc':
bias = not args.no_bias # if no_bias is True then bias must be False
net = fc(bias=bias,
width=args.width,
depth=args.depth,
input_dim=input_dim,
num_classes=num_classes).to(device)
else:
raise NotImplementedError
print(net)
torch.manual_seed(args.seed)
if args.zero_init_method == 'centered':
with torch.no_grad():
out0_list = []
for x, y, idx in train_loader_eval:
out0_list.append(net(x))
out0 = torch.cat(out0_list)
evaluation_history_TEST = torch.load(args.load_dir +
'/evaluation_history_TEST.hist')
evaluation_history_TRAIN = torch.load(args.load_dir +
'/evaluation_history_TRAIN.hist')
noise_norm_history_TEST = torch.load(args.load_dir +
'/noise_norm_history_TEST.hist')
noise_norm_history_TRAIN = torch.load(args.load_dir +
'/noise_norm_history_TRAIN.hist')
else:
if args.model == 'fc':
if args.dataset == 'mnist' or args.dataset == 'gauss':
if args.dataset == 'gauss':
net = fc(width=args.width,
depth=args.depth,
num_classes=num_classes,
input_dim=args.gauss_dimension,
activation=args.activation,
use_batch_norm=args.batch_norm).to(args.device)
else:
net = fc(width=args.width,
depth=args.depth,
num_classes=num_classes,
activation=args.activation,
use_batch_norm=args.batch_norm).to(args.device)
elif args.dataset == 'cifar10':
net = fc(width=args.width,
depth=args.depth,
num_classes=num_classes,
input_dim=3 * 32 * 32,
activation=args.activation,
use_batch_norm=args.batch_norm).to(args.device)