2021SC@SDUSC
在最后,我们对于train.py文件进行分析,此文件的作用是为了对TMix模型和MixText模型进行训练。
首先第一段代码与之前的类似,则不做过多赘述
parser = argparse.ArgumentParser(description='PyTorch MixText')
parser.add_argument('--epochs', default=50, type=int, metavar='N',
help='number of total epochs to run')
parser.add_argument('--batch-size', default=4, type=int, metavar='N',
help='train batchsize')
parser.add_argument('--batch-size-u', default=24, type=int, metavar='N',
help='train batchsize')
parser.add_argument('--lrmain', '--learning-rate-bert', default=0.00001, type=float,
metavar='LR', help='initial learning rate for bert')
parser.add_argument('--lrlast', '--learning-rate-model', default=0.001, type=float,
metavar='LR', help='initial learning rate for models')
parser.add_argument('--gpu', default='0,1,2,3', type=str,
help='id(s) for CUDA_VISIBLE_DEVICES')
parser.add_argument('--n-labeled', type=int, default=20,
help='number of labeled data')
parser.add_argument('--un-labeled', default=5000, type=int,
help='number of unlabeled data')
parser.add_argument('--val-iteration', type=int, default=200,
help='number of labeled data')
parser.add_argument('--mix-option', default=True, type=bool, metavar='N',
help='mix option, whether to mix or not')
parser.add_argument('--mix-method', default=0, type=int, metavar='N',
help='mix method, set different mix method')
parser.add_argument('--separate-mix', default=False, type=bool, metavar='N',
help='mix separate from labeled data and unlabeled data')
parser.add_argument('--co', default=False, type=bool, metavar='N',
help='set a random choice between mix and unmix during training')
parser.add_argument('--train_aug', default=False, type=bool, metavar='N',
help='augment labeled training data')
parser.add_argument('--model', type=str, default='bert-base-uncased',
help='pretrained model')
parser.add_argument('--data-path', type=str, default='/Users/wuzehao/Desktop/科研/文本分类/MixText-master/data/yahoo_answers_csv/',
help='path to data folders')
parser.add_argument('--mix-layers-set', nargs='+',
default=[0, 1, 2, 3], type=int, help='define mix layer set')
parser.add_argument('--alpha', default=0.75, type=float,
help='alpha for beta distribution')
parser.add_argument('--lambda-u', default=1, type=float,
help='weight for consistency loss term of unlabeled data')
parser.add_argument('--T', default=0.5, type=float,
help='temperature for sharpen function')
parser.add_argument('--temp-change', default=1000000, type=int)
parser.add_argument('--margin', default=0.7, type=float, metavar='N',
help='margin for hinge loss')
parser.add_argument('--lambda-u-hinge', default=0, type=float,
help='weight for hinge loss term of unlabeled data')
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
n_gpu = torch.cuda.device_count()
print("GPU num: ", n_gpu)
best_acc = 0
total_steps = 0
flag = 0
print('Whether mix: ', args.mix_option)
print("Mix layers sets: ", args.mix_layers_set)
我们首先对于train方法进行分析:
首先是对于未标记数据进行标签预测:
with torch.no_grad():
outputs_u = model(inputs_u)
outputs_u2 = model(inputs_u2)
outputs_ori = model(inputs_ori)
根据翻译质量,可以选择不同的权重
对于AG News:German为1,Russian:0,ori:1
对于DBPedia:German为1,Russian:1,ori:1
对于IMDB:German为0,Russian:0,ori:1
对于Yahoo Answers:German为1,Russian:0,ori:1或者erman为0,Russian:0,ori:1
if args.mix_method == 0:
logits = model(input_a, input_b, l, mix_layer)
mixed_target = l * target_a + (1 - l) * target_b
上方代码是混合句子的隐藏表示
elif args.mix_method == 1:
mixed_input = []
if l != 1:
for i in range(input_a.size(0)):
length1 = math.floor(int(length_a[i]) * l)
idx1 = torch.randperm(int(length_a[i]) - length1 + 1)[0]
length2 = math.ceil(int(length_b[i]) * (1-l))
if length1 + length2 > 256:
length2 = 256-length1 - 1
idx2 = torch.randperm(int(length_b[i]) - length2 + 1)[0]
try:
mixed_input.append(
torch.cat((input_a[i][idx1: idx1 + length1], torch.tensor([102]).cuda(), input_b[i][idx2:idx2 + length2], torch.tensor([0]*(256-1-length1-length2)).cuda()), dim=0).unsqueeze(0))
except:
print(256 - 1 - length1 - length2,
idx2, length2, idx1, length1)
mixed_input = torch.cat(mixed_input, dim=0)
else:
mixed_input = input_a
logits = model(mixed_input)
mixed_target = l * target_a + (1 - l) * target_b
上述代码实现了连接两个训练语句片段,然后这些片段是根据lamda进行选择的。例如:“I love you so much"和“He likes NLP”可以被混合成“He likes NLP so much”。相关联的标签会根据系数进行混合。
elif args.mix_method == 2:
if l == 1:
mixed_input = []
for i in range(input_a.size(0)):
mixed_input.append(
torch.cat((input_a[i][:length_a[i]], torch.tensor([102]).cuda(), input_b[i][:length_b[i]], torch.tensor([0]*(512-1-int(length_a[i])-int(length_b[i]))).cuda()), dim=0).unsqueeze(0))
mixed_input = torch.cat(mixed_input, dim=0)
logits = model(mixed_input, sent_size=512)
上述代码会合并训练语句。