def get_models(args, train=True, as_ensemble=False, model_file=None, leaky_relu=False):
models = []
mean = torch.tensor([0.4914, 0.4822, 0.4465], dtype=torch.float32).cuda()
std = torch.tensor([0.2023, 0.1994, 0.2010], dtype=torch.float32).cuda()
normalizer = NormalizeByChannelMeanStd(mean=mean, std=std)
if model_file:
state_dict = torch.load(model_file)
if train:
print('Loading pre-trained models...')
iter_m = state_dict.keys() if model_file else range(args.model_num)
for i in iter_m:
if args.arch.lower() == 'resnet':
model = ResNet(depth=args.depth, leaky_relu=leaky_relu)
else:
raise ValueError('[{:s}] architecture is not supported yet...')
# we include input normalization as a part of the model
model = ModelWrapper(model, normalizer)
if model_file:
model.load_state_dict(state_dict[i])
if train:
model.train()
else:
model.eval()
model = model.cuda()
models.append(model)
if as_ensemble:
assert not train, 'Must be in eval mode when getting models to form an ensemble'
ensemble = Ensemble(models)
ensemble.eval()
return ensemble
else:
return models
pin_memory=True,
sampler=None)
assert pred['wnids'][:1000] == train_wnids
model = ResNet('resnet50', 1000)
sd = model.resnet_base.state_dict()
sd.update(torch.load('materials/resnet50-base.pth'))
model.resnet_base.load_state_dict(sd)
fcw = pred['pred'][:1000].cpu()
model.fc.weight = nn.Parameter(fcw[:, :-1])
model.fc.bias = nn.Parameter(fcw[:, -1])
model = model.cuda()
model.train()
optimizer = torch.optim.SGD(model.resnet_base.parameters(),
lr=0.0001,
momentum=0.9)
loss_fn = nn.CrossEntropyLoss().cuda()
keep_ratio = 0.9975
trlog = {}
trlog['loss'] = []
trlog['acc'] = []
for epoch in range(1, 9999):
ave_loss = None
ave_acc = None
