###############################################################################
# Load data
###############################################################################
print('==> Preparing data..')
trainloader, validloader, testloader, nb_classes, dim_inp = get_dataset(args)
###############################################################################
# Build the model
###############################################################################
epoch = 0
if args.load_dir == '':
inp_channels = 3
print('==> Building model..')
if args.arch == 'resnet':
model0 = ResNet_model(bn= args.bn, num_classes=nb_classes, depth=args.depth,\
inp_channels=inp_channels, k=args.k, affine=not args.noaffine, inp_noise=args.inp_noise, VIB=args.vib)
elif args.arch == 'cnn':
model0 = CNN(bn=args.bn,
affine=not args.noaffine,
num_classes=nb_classes,
inp_noise=args.inp_noise,
VIB=args.vib)
else:
with open(args.root_dir + '/' + args.load_dir + '/best_model.pt',
'rb') as f:
best_state = torch.load(f)
model0 = best_state['model']
epoch = best_state['epoch']
print('==> Loading model from epoch ', epoch)
params = list(model0.parameters())
testloader = torch.utils.data.DataLoader(testset, batch_size=args.mbs, shuffle=False,
num_workers=0, pin_memory=True)
nb_classes = 10
dim_inp = 32 * 32 * 3
###############################################################################
# Build the model
###############################################################################
inp_channels = 1 if args.dataset=='mnist' else 3
print('==> Building model..')
start_epoch=0
if args.arch == 'linear':
model = LinearNet(dim_inp, nb_classes)
elif args.arch == 'resnet':
model = ResNet_model(dropout=args.dropout, normalization= args.normalization, num_classes=nb_classes, dataset=args.dataset, depth=args.depth, nb_filters=args.nb_filters, kernel_size=args.kernel,\
inp_channels=inp_channels, k=args.k, affine=not args.noaffine)
params = list(model.parameters())
model = torch.nn.DataParallel(model, device_ids=range(len(args.gpu)))
nb = 0
if args.init == 'he':
for m in model.modules():
if isinstance(m, nn.Conv2d):
nb += 1
print ('Update init of ', m)
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d) and not args.noaffine:
print ('Update init of ', m)