def get_model(args):
sd = None
model_args = args
if args.load is not None and args.load != '':
# sd = torch.load(args.load, map_location=lambda storage, location: 'cpu')
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
sd = torch.load(args.load, map_location=device)
if 'args' in sd:
model_args = sd['args']
if 'sd' in sd:
sd = sd['sd']
ntokens = model_args.data_size
concat_pools = model_args.concat_max, model_args.concat_min, model_args.concat_mean
if args.model == 'transformer':
model = SentimentClassifier(model_args.model, ntokens, None, None,
None, model_args.classifier_hidden_layers,
model_args.classifier_dropout, None,
concat_pools, False, model_args)
else:
model = SentimentClassifier(
model_args.model, ntokens, model_args.emsize, model_args.nhid,
model_args.nlayers, model_args.classifier_hidden_layers,
model_args.classifier_dropout, model_args.all_layers, concat_pools,
False, model_args)
args.heads_per_class = model_args.heads_per_class
args.use_softmax = model_args.use_softmax
try:
args.classes = list(model_args.classes)
except:
args.classes = [args.label_key]
try:
args.dual_thresh = model_args.dual_thresh and not model_args.joint_binary_train
except:
args.dual_thresh = False
if args.cuda:
model.cuda()
if args.fp16:
model.half()
if sd is not None:
try:
model.load_state_dict(sd)
except:
# if state dict has weight normalized parameters apply and remove weight norm to model while loading sd
if hasattr(model.lm_encoder, 'rnn'):
apply_weight_norm(model.lm_encoder.rnn)
else:
apply_weight_norm(model.lm_encoder)
model.lm_encoder.load_state_dict(sd)
remove_weight_norm(model)
if args.neurons > 0:
print('WARNING. Setting neurons %s' % str(args.neurons))
model.set_neurons(args.neurons)
return model
(key, old_args[key], key, old_args[key]))
old_args[key] = args[key]
# replace args
args = old_args
# clean up a bit
args.pop("checkpoint", None)
print("New args:")
print(args)
# Instantiating the classifier model
print("Building Model")
model = SentimentClassifier(args["freeze_bert"], args["dropout"])
model.cuda(args["gpu"]) # Enable gpu support for the model
print("Creating criterion and optimizer objects")
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.SGD(model.parameters(), lr=args["lr"])
# Creating dataloaders
print("Creating train and val dataloaders")
dataset_train = SSTDataset(filename='data/SST-2/train.tsv',
maxlen=args["maxlen"])
dataset_validation = SSTDataset(filename='data/SST-2/dev.tsv',
maxlen=args["maxlen"])
train_loader = DataLoader(dataset_train,
batch_size=args["batch_size"],
num_workers=5,
help='number of nenurons to extract as features')
data_config, data_parser = configure_data(parser)
data_parser.set_defaults(split='1.', data='data/binary_sst/train.csv')
args = parser.parse_args()
args.cuda = torch.cuda.is_available()
train_data, val_data, test_data = data_config.apply(args)
ntokens = args.data_size
model = SentimentClassifier(args.model, ntokens, args.emsize, args.nhid,
args.nlayers, 0.0, args.all_layers)
if args.cuda:
model.cuda()
if args.fp16:
model.half()
with open(args.load_model, 'rb') as f:
sd = torch.load(f)
try:
model.load_state_dict(sd)
except:
apply_weight_norm(model.encoder.rnn)
model.load_state_dict(sd)
remove_weight_norm(model)
if args.neurons > 0: