mode='train',
model_type='rnn'),
batch_size=args.train_batch,
num_workers=args.workers,
shuffle=True)
# collate_fn=data.collate_fn)
val_loader = torch.utils.data.DataLoader(data.Data(args,
mode='val',
model_type='rnn'),
batch_size=args.train_batch,
num_workers=args.workers,
shuffle=False)
''' load model '''
print('===> prepare model ...')
feature_extractor, RNNClassifier = model.Resnet50(), model.RNNClassifier(
args)
feature_extractor.cuda()
RNNClassifier.cuda()
clf_criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(RNNClassifier.parameters(),
lr=args.lr,
momentum=0.9)
''' setup tensorboard '''
writer = SummaryWriter(os.path.join(args.save_dir, 'train_info'))
print('===> start training ...')
iters = 0
best_acc = 0
for epoch in range(1, args.epoch + 1):
test_data = data.sentimentDataset(vocab, csv_dataset, train=False)
train_dataloader = DataLoader(train_data,
batch_size=config.TRAIN_BATCH_SIZE,
shuffle=True,
collate_fn=data.collate_fn)
test_dataloader = DataLoader(test_data,
batch_size=config.TEST_BATCH_SIZE,
shuffle=True,
collate_fn=data.collate_fn)
model_classifier = model.RNNClassifier(nembedding=config.DIM,
hidden_size=config.HIDDEN_SIZE,
num_layer=config.NUM_LAYER,
dropout=config.drop_out,
vocab_size=vocab.n_words,
use_pretrain=False,
embed_matrix=vocab.vector,
embed_freeze=False,
label_size=2)
model_classifier = model.CNNClassifier(nembedding=config.DIM,
vocab_size=vocab.n_words,
kernel_num=3,
kernel_sizes=config.kernel_sizes,
label_size=2,
dropout=0.3,
use_pretrain=False)
use_loaded_model = input("load model?")
if use_loaded_model == "yes":
args.save_dir = os.path.join(
args.save_dir,
datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S'))
print("\nParameters:")
for attr, value in sorted(args.__dict__.items()):
print("\t{}={}".format(attr.upper(), value))
# model
if args.snapshot is not None:
print('\nLoading model from {}...'.format(args.snapshot))
classification_model = torch.load(args.snapshot)
elif args.rnn:
classification_model = model.RNNClassifier(args.embed_num, args.embed_num,
args.hidden_dim,
args.rnn_layers, args.class_num)
else:
classification_model = model.CharCNN(args.embed_num, args.embed_dim,
args.class_num, args.kernel_num,
args.kernel_sizes, args.dropout)
if args.cuda:
torch.cuda.set_device(args.device)
classification_model = classification_model.cuda()
# train or predict
if args.predict is not None:
label = predict.predict(args.predict, classification_model, text_field,
label_field, args.cuda)
print('\n[Text] {}\n[Label] {}\n'.format(args.predict, label))