#!/usr/bin/python # Copyright (c) Facebook, Inc. and its affiliates. # All rights reserved. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. # # LASER Language-Agnostic SEntence Representations # is a toolkit to calculate multilingual sentence embeddings # and to use them for document classification, bitext filtering # and mining # # -------------------------------------------------------- # # Simple MLP classifier for sentence embeddings import argparse import copy import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim import torch.utils.data as data_utils ################################################ def LoadData(bdir, dfn, lfn, dim=1024, bsize=32, shuffle=False, quiet=False): x = np.fromfile(bdir + dfn, dtype=np.float32, count=-1) x.resize(x.shape[0] // dim, dim) lbl = np.loadtxt(bdir + lfn, dtype=np.int32) lbl.reshape(lbl.shape[0], 1) if not quiet: print(' - read {:d}x{:d} elements in {:s}'.format(x.shape[0], x.shape[1], dfn)) print(' - read {:d} labels [{:d},{:d}] in {:s}' .format(lbl.shape[0], lbl.min(), lbl.max(), lfn)) D = data_utils.TensorDataset(torch.from_numpy(x), torch.from_numpy(lbl)) loader = data_utils.DataLoader(D, batch_size=bsize, shuffle=shuffle) return loader ################################################ class Net(nn.Module): def __init__(self, idim=1024, odim=2, nhid=None, dropout=0.0, gpu=0, activation='TANH'): super(Net, self).__init__() self.gpu = gpu modules = [] modules = [] print(' - mlp {:d}'.format(idim), end='') if len(nhid) > 0: if dropout > 0: modules.append(nn.Dropout(p=dropout)) nprev = idim for nh in nhid: if nh > 0: modules.append(nn.Linear(nprev, nh)) nprev = nh if activation == 'TANH': modules.append(nn.Tanh()) print('-{:d}t'.format(nh), end='') elif activation == 'RELU': modules.append(nn.ReLU()) print('-{:d}r'.format(nh), end='') else: raise Exception('Unrecognized activation {activation}') if dropout > 0: modules.append(nn.Dropout(p=dropout)) modules.append(nn.Linear(nprev, odim)) print('-{:d}, dropout={:.1f}'.format(odim, dropout)) else: modules.append(nn.Linear(idim, odim)) print(' - mlp %d-%d'.format(idim, odim)) self.mlp = nn.Sequential(*modules) # Softmax is included CrossEntropyLoss ! if self.gpu >= 0: self.mlp = self.mlp.cuda() def forward(self, x): return self.mlp(x) def TestCorpus(self, dset, name=' Dev', nlbl=4): correct = 0 total = 0 self.mlp.train(mode=False) corr = np.zeros(nlbl, dtype=np.int32) for data in dset: X, Y = data Y = Y.long() if self.gpu >= 0: X = X.cuda() Y = Y.cuda() outputs = self.mlp(X) _, predicted = torch.max(outputs.data, 1) total += Y.size(0) correct += (predicted == Y).int().sum() for i in range(nlbl): corr[i] += (predicted == i).int().sum() print(' | {:4s}: {:5.2f}%' .format(name, 100.0 * correct.float() / total), end='') print(' | classes:', end='') for i in range(nlbl): print(' {:5.2f}'.format(100.0 * corr[i] / total), end='') return correct, total ################################################ parser = argparse.ArgumentParser( formatter_class=argparse.RawDescriptionHelpFormatter, description="Simple sentence classifier") # Data parser.add_argument( '--base-dir', '-b', type=str, required=True, metavar='PATH', help="Directory with all the data files)") parser.add_argument( '--save', '-s', type=str, required=False, metavar='PATH', default="", help="File in which to save best network") parser.add_argument( '--train', '-t', type=str, required=True, metavar='STR', help="Name of training corpus") parser.add_argument( '--train-labels', '-T', type=str, required=True, metavar='STR', help="Name of training corpus (labels)") parser.add_argument( '--dev', '-d', type=str, required=True, metavar='STR', help="Name of development corpus") parser.add_argument( '--dev-labels', '-D', type=str, required=True, metavar='STR', help="Name of development corpus (labels)") parser.add_argument( '--test', '-e', type=str, required=True, metavar='STR', help="Name of test corpus without language extension") parser.add_argument( '--test-labels', '-E', type=str, required=True, metavar='STR', help="Name of test corpus without language extension (labels)") parser.add_argument( '--lang', '-L', nargs='+', default=None, help="List of languages to test on") # network definition parser.add_argument( "--dim", "-m", type=int, default=1024, help="Dimension of sentence embeddings") parser.add_argument( '--nhid', '-n', type=int, default=[0], nargs='+', help="List of hidden layer(s) dimensions") parser.add_argument( "--nb-classes", "-c", type=int, default=2, help="Number of output classes") parser.add_argument( '--dropout', '-o', type=float, default=0.0, metavar='FLOAT', help="Value of dropout") parser.add_argument( '--nepoch', '-N', type=int, default=100, metavar='INT', help="Number of epochs") parser.add_argument( '--bsize', '-B', type=int, default=128, metavar='INT', help="Batch size") parser.add_argument( '--seed', '-S', type=int, default=123456789, metavar='INT', help="Initial random seed") parser.add_argument( '--lr', type=float, default=0.001, metavar='FLOAT', help='Learning rate') parser.add_argument( '--wdecay', type=float, default=0.0, metavar='FLOAT', help='Weight decay') parser.add_argument( '--gpu', '-g', type=int, default=-1, metavar='INT', help="GPU id (-1 for CPU)") args = parser.parse_args() print(' - base directory: {}'.format(args.base_dir)) args.base_dir = args.base_dir + "/" train_loader = LoadData(args.base_dir, args.train, args.train_labels, dim=args.dim, bsize=args.bsize, shuffle=True) dev_loader = LoadData(args.base_dir, args.dev, args.dev_labels, dim=args.dim, bsize=args.bsize, shuffle=False) # set GPU and random seed torch.cuda.set_device(args.gpu) np.random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) print(" - setting seed to %d" % args.seed) # create network net = Net(idim=args.dim, odim=args.nb_classes, nhid=args.nhid, dropout=args.dropout, gpu=args.gpu) if args.gpu >= 0: criterion = nn.CrossEntropyLoss().cuda() else: criterion = nn.CrossEntropyLoss() #optimizer = optim.Adam(net.parameters(), weight_decay=0.0) # default: pytorch/optim/adam.py # Py0.4: lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, amsgrad=False): # Py1.0: lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, amsgrad=False): optimizer = optim.Adam(net.parameters(), lr=args.lr, weight_decay=args.wdecay, betas=(0.9, 0.999), eps=1e-8, amsgrad=False) corr_best = 0 # loop multiple times over the dataset for epoch in range(args.nepoch): loss_epoch = 0.0 print('Ep {:4d}'.format(epoch), end='') # for inputs, labels in train_loader: for i, data in enumerate(train_loader, 0): # get the inputs inputs, labels = data labels = labels.long() if args.gpu >= 0: inputs = inputs.cuda() labels = labels.cuda() # zero the parameter gradients net.zero_grad() # forward + backward + optimize net.train(mode=True) outputs = net(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() loss_epoch += loss.item() print(' | loss {:e}'.format(loss_epoch), end='') corr, nbex = net.TestCorpus(dev_loader, 'Dev') if corr >= corr_best: print(' | saved') corr_best = corr net_best = copy.deepcopy(net) else: print('') if 'net_best' in globals(): if args.save != '': torch.save(net_best.cpu(), args.save) print('Best Dev: {:d} = {:5.2f}%' .format(corr_best, 100.0 * corr_best.float() / nbex)) if args.gpu >= 0: net_best = net_best.cuda() # test on (several) languages for l in args.lang: test_loader = LoadData(args.base_dir, args.test + '.' + l, args.test_labels + '.' + l, dim=args.dim, bsize=args.bsize, shuffle=False, quiet=True) print('Ep best | Eval Test lang {:s}'.format(l), end='') net_best.TestCorpus(test_loader, 'Test') print('')