@@ -0,0 +1,223 @@ | |||
# Code Modified from https://github.com/carpedm20/ENAS-pytorch | |||
"""A module with NAS controller-related code.""" | |||
import collections | |||
import os | |||
import torch | |||
import torch.nn.functional as F | |||
import fastNLP | |||
import fastNLP.models.enas_utils as utils | |||
from fastNLP.models.enas_utils import Node | |||
def _construct_dags(prev_nodes, activations, func_names, num_blocks): | |||
"""Constructs a set of DAGs based on the actions, i.e., previous nodes and | |||
activation functions, sampled from the controller/policy pi. | |||
Args: | |||
prev_nodes: Previous node actions from the policy. | |||
activations: Activations sampled from the policy. | |||
func_names: Mapping from activation function names to functions. | |||
num_blocks: Number of blocks in the target RNN cell. | |||
Returns: | |||
A list of DAGs defined by the inputs. | |||
RNN cell DAGs are represented in the following way: | |||
1. Each element (node) in a DAG is a list of `Node`s. | |||
2. The `Node`s in the list dag[i] correspond to the subsequent nodes | |||
that take the output from node i as their own input. | |||
3. dag[-1] is the node that takes input from x^{(t)} and h^{(t - 1)}. | |||
dag[-1] always feeds dag[0]. | |||
dag[-1] acts as if `w_xc`, `w_hc`, `w_xh` and `w_hh` are its | |||
weights. | |||
4. dag[N - 1] is the node that produces the hidden state passed to | |||
the next timestep. dag[N - 1] is also always a leaf node, and therefore | |||
is always averaged with the other leaf nodes and fed to the output | |||
decoder. | |||
""" | |||
dags = [] | |||
for nodes, func_ids in zip(prev_nodes, activations): | |||
dag = collections.defaultdict(list) | |||
# add first node | |||
dag[-1] = [Node(0, func_names[func_ids[0]])] | |||
dag[-2] = [Node(0, func_names[func_ids[0]])] | |||
# add following nodes | |||
for jdx, (idx, func_id) in enumerate(zip(nodes, func_ids[1:])): | |||
dag[utils.to_item(idx)].append(Node(jdx + 1, func_names[func_id])) | |||
leaf_nodes = set(range(num_blocks)) - dag.keys() | |||
# merge with avg | |||
for idx in leaf_nodes: | |||
dag[idx] = [Node(num_blocks, 'avg')] | |||
# This is actually y^{(t)}. h^{(t)} is node N - 1 in | |||
# the graph, where N Is the number of nodes. I.e., h^{(t)} takes | |||
# only one other node as its input. | |||
# last h[t] node | |||
last_node = Node(num_blocks + 1, 'h[t]') | |||
dag[num_blocks] = [last_node] | |||
dags.append(dag) | |||
return dags | |||
class Controller(torch.nn.Module): | |||
"""Based on | |||
https://github.com/pytorch/examples/blob/master/word_language_model/model.py | |||
RL controllers do not necessarily have much to do with | |||
language models. | |||
Base the controller RNN on the GRU from: | |||
https://github.com/ikostrikov/pytorch-a2c-ppo-acktr/blob/master/model.py | |||
""" | |||
def __init__(self, num_blocks=4, controller_hid=100, cuda=False): | |||
torch.nn.Module.__init__(self) | |||
# `num_tokens` here is just the activation function | |||
# for every even step, | |||
self.shared_rnn_activations = ['tanh', 'ReLU', 'identity', 'sigmoid'] | |||
self.num_tokens = [len(self.shared_rnn_activations)] | |||
self.controller_hid = controller_hid | |||
self.use_cuda = cuda | |||
self.num_blocks = num_blocks | |||
for idx in range(num_blocks): | |||
self.num_tokens += [idx + 1, len(self.shared_rnn_activations)] | |||
self.func_names = self.shared_rnn_activations | |||
num_total_tokens = sum(self.num_tokens) | |||
self.encoder = torch.nn.Embedding(num_total_tokens, | |||
controller_hid) | |||
self.lstm = torch.nn.LSTMCell(controller_hid, controller_hid) | |||
# Perhaps these weights in the decoder should be | |||
# shared? At least for the activation functions, which all have the | |||
# same size. | |||
self.decoders = [] | |||
for idx, size in enumerate(self.num_tokens): | |||
decoder = torch.nn.Linear(controller_hid, size) | |||
self.decoders.append(decoder) | |||
self._decoders = torch.nn.ModuleList(self.decoders) | |||
self.reset_parameters() | |||
self.static_init_hidden = utils.keydefaultdict(self.init_hidden) | |||
def _get_default_hidden(key): | |||
return utils.get_variable( | |||
torch.zeros(key, self.controller_hid), | |||
self.use_cuda, | |||
requires_grad=False) | |||
self.static_inputs = utils.keydefaultdict(_get_default_hidden) | |||
def reset_parameters(self): | |||
init_range = 0.1 | |||
for param in self.parameters(): | |||
param.data.uniform_(-init_range, init_range) | |||
for decoder in self.decoders: | |||
decoder.bias.data.fill_(0) | |||
def forward(self, # pylint:disable=arguments-differ | |||
inputs, | |||
hidden, | |||
block_idx, | |||
is_embed): | |||
if not is_embed: | |||
embed = self.encoder(inputs) | |||
else: | |||
embed = inputs | |||
hx, cx = self.lstm(embed, hidden) | |||
logits = self.decoders[block_idx](hx) | |||
logits /= 5.0 | |||
# # exploration | |||
# if self.args.mode == 'train': | |||
# logits = (2.5 * F.tanh(logits)) | |||
return logits, (hx, cx) | |||
def sample(self, batch_size=1, with_details=False, save_dir=None): | |||
"""Samples a set of `args.num_blocks` many computational nodes from the | |||
controller, where each node is made up of an activation function, and | |||
each node except the last also includes a previous node. | |||
""" | |||
if batch_size < 1: | |||
raise Exception(f'Wrong batch_size: {batch_size} < 1') | |||
# [B, L, H] | |||
inputs = self.static_inputs[batch_size] | |||
hidden = self.static_init_hidden[batch_size] | |||
activations = [] | |||
entropies = [] | |||
log_probs = [] | |||
prev_nodes = [] | |||
# The RNN controller alternately outputs an activation, | |||
# followed by a previous node, for each block except the last one, | |||
# which only gets an activation function. The last node is the output | |||
# node, and its previous node is the average of all leaf nodes. | |||
for block_idx in range(2*(self.num_blocks - 1) + 1): | |||
logits, hidden = self.forward(inputs, | |||
hidden, | |||
block_idx, | |||
is_embed=(block_idx == 0)) | |||
probs = F.softmax(logits, dim=-1) | |||
log_prob = F.log_softmax(logits, dim=-1) | |||
# .mean() for entropy? | |||
entropy = -(log_prob * probs).sum(1, keepdim=False) | |||
action = probs.multinomial(num_samples=1).data | |||
selected_log_prob = log_prob.gather( | |||
1, utils.get_variable(action, requires_grad=False)) | |||
# why the [:, 0] here? Should it be .squeeze(), or | |||
# .view()? Same below with `action`. | |||
entropies.append(entropy) | |||
log_probs.append(selected_log_prob[:, 0]) | |||
# 0: function, 1: previous node | |||
mode = block_idx % 2 | |||
inputs = utils.get_variable( | |||
action[:, 0] + sum(self.num_tokens[:mode]), | |||
requires_grad=False) | |||
if mode == 0: | |||
activations.append(action[:, 0]) | |||
elif mode == 1: | |||
prev_nodes.append(action[:, 0]) | |||
prev_nodes = torch.stack(prev_nodes).transpose(0, 1) | |||
activations = torch.stack(activations).transpose(0, 1) | |||
dags = _construct_dags(prev_nodes, | |||
activations, | |||
self.func_names, | |||
self.num_blocks) | |||
if save_dir is not None: | |||
for idx, dag in enumerate(dags): | |||
utils.draw_network(dag, | |||
os.path.join(save_dir, f'graph{idx}.png')) | |||
if with_details: | |||
return dags, torch.cat(log_probs), torch.cat(entropies) | |||
return dags | |||
def init_hidden(self, batch_size): | |||
zeros = torch.zeros(batch_size, self.controller_hid) | |||
return (utils.get_variable(zeros, self.use_cuda, requires_grad=False), | |||
utils.get_variable(zeros.clone(), self.use_cuda, requires_grad=False)) |
@@ -0,0 +1,388 @@ | |||
# Code Modified from https://github.com/carpedm20/ENAS-pytorch | |||
"""Module containing the shared RNN model.""" | |||
import numpy as np | |||
import collections | |||
import torch | |||
from torch import nn | |||
import torch.nn.functional as F | |||
from torch.autograd import Variable | |||
import fastNLP.models.enas_utils as utils | |||
from fastNLP.models.base_model import BaseModel | |||
import fastNLP.modules.encoder as encoder | |||
def _get_dropped_weights(w_raw, dropout_p, is_training): | |||
"""Drops out weights to implement DropConnect. | |||
Args: | |||
w_raw: Full, pre-dropout, weights to be dropped out. | |||
dropout_p: Proportion of weights to drop out. | |||
is_training: True iff _shared_ model is training. | |||
Returns: | |||
The dropped weights. | |||
Why does torch.nn.functional.dropout() return: | |||
1. `torch.autograd.Variable()` on the training loop | |||
2. `torch.nn.Parameter()` on the controller or eval loop, when | |||
training = False... | |||
Even though the call to `_setweights` in the Smerity repo's | |||
`weight_drop.py` does not have this behaviour, and `F.dropout` always | |||
returns `torch.autograd.Variable` there, even when `training=False`? | |||
The above TODO is the reason for the hacky check for `torch.nn.Parameter`. | |||
""" | |||
dropped_w = F.dropout(w_raw, p=dropout_p, training=is_training) | |||
if isinstance(dropped_w, torch.nn.Parameter): | |||
dropped_w = dropped_w.clone() | |||
return dropped_w | |||
class EmbeddingDropout(torch.nn.Embedding): | |||
"""Class for dropping out embeddings by zero'ing out parameters in the | |||
embedding matrix. | |||
This is equivalent to dropping out particular words, e.g., in the sentence | |||
'the quick brown fox jumps over the lazy dog', dropping out 'the' would | |||
lead to the sentence '### quick brown fox jumps over ### lazy dog' (in the | |||
embedding vector space). | |||
See 'A Theoretically Grounded Application of Dropout in Recurrent Neural | |||
Networks', (Gal and Ghahramani, 2016). | |||
""" | |||
def __init__(self, | |||
num_embeddings, | |||
embedding_dim, | |||
max_norm=None, | |||
norm_type=2, | |||
scale_grad_by_freq=False, | |||
sparse=False, | |||
dropout=0.1, | |||
scale=None): | |||
"""Embedding constructor. | |||
Args: | |||
dropout: Dropout probability. | |||
scale: Used to scale parameters of embedding weight matrix that are | |||
not dropped out. Note that this is _in addition_ to the | |||
`1/(1 - dropout)` scaling. | |||
See `torch.nn.Embedding` for remaining arguments. | |||
""" | |||
torch.nn.Embedding.__init__(self, | |||
num_embeddings=num_embeddings, | |||
embedding_dim=embedding_dim, | |||
max_norm=max_norm, | |||
norm_type=norm_type, | |||
scale_grad_by_freq=scale_grad_by_freq, | |||
sparse=sparse) | |||
self.dropout = dropout | |||
assert (dropout >= 0.0) and (dropout < 1.0), ('Dropout must be >= 0.0 ' | |||
'and < 1.0') | |||
self.scale = scale | |||
def forward(self, inputs): # pylint:disable=arguments-differ | |||
"""Embeds `inputs` with the dropped out embedding weight matrix.""" | |||
if self.training: | |||
dropout = self.dropout | |||
else: | |||
dropout = 0 | |||
if dropout: | |||
mask = self.weight.data.new(self.weight.size(0), 1) | |||
mask.bernoulli_(1 - dropout) | |||
mask = mask.expand_as(self.weight) | |||
mask = mask / (1 - dropout) | |||
masked_weight = self.weight * Variable(mask) | |||
else: | |||
masked_weight = self.weight | |||
if self.scale and self.scale != 1: | |||
masked_weight = masked_weight * self.scale | |||
return F.embedding(inputs, | |||
masked_weight, | |||
max_norm=self.max_norm, | |||
norm_type=self.norm_type, | |||
scale_grad_by_freq=self.scale_grad_by_freq, | |||
sparse=self.sparse) | |||
class LockedDropout(nn.Module): | |||
# code from https://github.com/salesforce/awd-lstm-lm/blob/master/locked_dropout.py | |||
def __init__(self): | |||
super().__init__() | |||
def forward(self, x, dropout=0.5): | |||
if not self.training or not dropout: | |||
return x | |||
m = x.data.new(1, x.size(1), x.size(2)).bernoulli_(1 - dropout) | |||
mask = Variable(m, requires_grad=False) / (1 - dropout) | |||
mask = mask.expand_as(x) | |||
return mask * x | |||
class ENASModel(BaseModel): | |||
"""Shared RNN model.""" | |||
def __init__(self, embed_num, num_classes, num_blocks=4, cuda=False, shared_hid=1000, shared_embed=1000): | |||
super(ENASModel, self).__init__() | |||
self.use_cuda = cuda | |||
self.shared_hid = shared_hid | |||
self.num_blocks = num_blocks | |||
self.decoder = nn.Linear(self.shared_hid, num_classes) | |||
self.encoder = EmbeddingDropout(embed_num, | |||
shared_embed, | |||
dropout=0.1) | |||
self.lockdrop = LockedDropout() | |||
self.dag = None | |||
# Tie weights | |||
# self.decoder.weight = self.encoder.weight | |||
# Since W^{x, c} and W^{h, c} are always summed, there | |||
# is no point duplicating their bias offset parameter. Likewise for | |||
# W^{x, h} and W^{h, h}. | |||
self.w_xc = nn.Linear(shared_embed, self.shared_hid) | |||
self.w_xh = nn.Linear(shared_embed, self.shared_hid) | |||
# The raw weights are stored here because the hidden-to-hidden weights | |||
# are weight dropped on the forward pass. | |||
self.w_hc_raw = torch.nn.Parameter( | |||
torch.Tensor(self.shared_hid, self.shared_hid)) | |||
self.w_hh_raw = torch.nn.Parameter( | |||
torch.Tensor(self.shared_hid, self.shared_hid)) | |||
self.w_hc = None | |||
self.w_hh = None | |||
self.w_h = collections.defaultdict(dict) | |||
self.w_c = collections.defaultdict(dict) | |||
for idx in range(self.num_blocks): | |||
for jdx in range(idx + 1, self.num_blocks): | |||
self.w_h[idx][jdx] = nn.Linear(self.shared_hid, | |||
self.shared_hid, | |||
bias=False) | |||
self.w_c[idx][jdx] = nn.Linear(self.shared_hid, | |||
self.shared_hid, | |||
bias=False) | |||
self._w_h = nn.ModuleList([self.w_h[idx][jdx] | |||
for idx in self.w_h | |||
for jdx in self.w_h[idx]]) | |||
self._w_c = nn.ModuleList([self.w_c[idx][jdx] | |||
for idx in self.w_c | |||
for jdx in self.w_c[idx]]) | |||
self.batch_norm = None | |||
# if args.mode == 'train': | |||
# self.batch_norm = nn.BatchNorm1d(self.shared_hid) | |||
# else: | |||
# self.batch_norm = None | |||
self.reset_parameters() | |||
self.static_init_hidden = utils.keydefaultdict(self.init_hidden) | |||
def setDAG(self, dag): | |||
if self.dag is None: | |||
self.dag = dag | |||
def forward(self, word_seq, hidden=None): | |||
inputs = torch.transpose(word_seq, 0, 1) | |||
time_steps = inputs.size(0) | |||
batch_size = inputs.size(1) | |||
self.w_hh = _get_dropped_weights(self.w_hh_raw, | |||
0.5, | |||
self.training) | |||
self.w_hc = _get_dropped_weights(self.w_hc_raw, | |||
0.5, | |||
self.training) | |||
# hidden = self.static_init_hidden[batch_size] if hidden is None else hidden | |||
hidden = self.static_init_hidden[batch_size] | |||
embed = self.encoder(inputs) | |||
embed = self.lockdrop(embed, 0.65 if self.training else 0) | |||
# The norm of hidden states are clipped here because | |||
# otherwise ENAS is especially prone to exploding activations on the | |||
# forward pass. This could probably be fixed in a more elegant way, but | |||
# it might be exposing a weakness in the ENAS algorithm as currently | |||
# proposed. | |||
# | |||
# For more details, see | |||
# https://github.com/carpedm20/ENAS-pytorch/issues/6 | |||
clipped_num = 0 | |||
max_clipped_norm = 0 | |||
h1tohT = [] | |||
logits = [] | |||
for step in range(time_steps): | |||
x_t = embed[step] | |||
logit, hidden = self.cell(x_t, hidden, self.dag) | |||
hidden_norms = hidden.norm(dim=-1) | |||
max_norm = 25.0 | |||
if hidden_norms.data.max() > max_norm: | |||
# Just directly use the torch slice operations | |||
# in PyTorch v0.4. | |||
# | |||
# This workaround for PyTorch v0.3.1 does everything in numpy, | |||
# because the PyTorch slicing and slice assignment is too | |||
# flaky. | |||
hidden_norms = hidden_norms.data.cpu().numpy() | |||
clipped_num += 1 | |||
if hidden_norms.max() > max_clipped_norm: | |||
max_clipped_norm = hidden_norms.max() | |||
clip_select = hidden_norms > max_norm | |||
clip_norms = hidden_norms[clip_select] | |||
mask = np.ones(hidden.size()) | |||
normalizer = max_norm/clip_norms | |||
normalizer = normalizer[:, np.newaxis] | |||
mask[clip_select] = normalizer | |||
if self.use_cuda: | |||
hidden *= torch.autograd.Variable( | |||
torch.FloatTensor(mask).cuda(), requires_grad=False) | |||
else: | |||
hidden *= torch.autograd.Variable( | |||
torch.FloatTensor(mask), requires_grad=False) | |||
logits.append(logit) | |||
h1tohT.append(hidden) | |||
h1tohT = torch.stack(h1tohT) | |||
output = torch.stack(logits) | |||
raw_output = output | |||
output = self.lockdrop(output, 0.4 if self.training else 0) | |||
#Pooling | |||
output = torch.mean(output, 0) | |||
decoded = self.decoder(output) | |||
extra_out = {'dropped': decoded, | |||
'hiddens': h1tohT, | |||
'raw': raw_output} | |||
return {'pred': decoded, 'hidden': hidden, 'extra_out': extra_out} | |||
def cell(self, x, h_prev, dag): | |||
"""Computes a single pass through the discovered RNN cell.""" | |||
c = {} | |||
h = {} | |||
f = {} | |||
f[0] = self.get_f(dag[-1][0].name) | |||
c[0] = torch.sigmoid(self.w_xc(x) + F.linear(h_prev, self.w_hc, None)) | |||
h[0] = (c[0]*f[0](self.w_xh(x) + F.linear(h_prev, self.w_hh, None)) + | |||
(1 - c[0])*h_prev) | |||
leaf_node_ids = [] | |||
q = collections.deque() | |||
q.append(0) | |||
# Computes connections from the parent nodes `node_id` | |||
# to their child nodes `next_id` recursively, skipping leaf nodes. A | |||
# leaf node is a node whose id == `self.num_blocks`. | |||
# | |||
# Connections between parent i and child j should be computed as | |||
# h_j = c_j*f_{ij}{(W^h_{ij}*h_i)} + (1 - c_j)*h_i, | |||
# where c_j = \sigmoid{(W^c_{ij}*h_i)} | |||
# | |||
# See Training details from Section 3.1 of the paper. | |||
# | |||
# The following algorithm does a breadth-first (since `q.popleft()` is | |||
# used) search over the nodes and computes all the hidden states. | |||
while True: | |||
if len(q) == 0: | |||
break | |||
node_id = q.popleft() | |||
nodes = dag[node_id] | |||
for next_node in nodes: | |||
next_id = next_node.id | |||
if next_id == self.num_blocks: | |||
leaf_node_ids.append(node_id) | |||
assert len(nodes) == 1, ('parent of leaf node should have ' | |||
'only one child') | |||
continue | |||
w_h = self.w_h[node_id][next_id] | |||
w_c = self.w_c[node_id][next_id] | |||
f[next_id] = self.get_f(next_node.name) | |||
c[next_id] = torch.sigmoid(w_c(h[node_id])) | |||
h[next_id] = (c[next_id]*f[next_id](w_h(h[node_id])) + | |||
(1 - c[next_id])*h[node_id]) | |||
q.append(next_id) | |||
# Instead of averaging loose ends, perhaps there should | |||
# be a set of separate unshared weights for each "loose" connection | |||
# between each node in a cell and the output. | |||
# | |||
# As it stands, all weights W^h_{ij} are doing double duty by | |||
# connecting both from i to j, as well as from i to the output. | |||
# average all the loose ends | |||
leaf_nodes = [h[node_id] for node_id in leaf_node_ids] | |||
output = torch.mean(torch.stack(leaf_nodes, 2), -1) | |||
# stabilizing the Updates of omega | |||
if self.batch_norm is not None: | |||
output = self.batch_norm(output) | |||
return output, h[self.num_blocks - 1] | |||
def init_hidden(self, batch_size): | |||
zeros = torch.zeros(batch_size, self.shared_hid) | |||
return utils.get_variable(zeros, self.use_cuda, requires_grad=False) | |||
def get_f(self, name): | |||
name = name.lower() | |||
if name == 'relu': | |||
f = torch.relu | |||
elif name == 'tanh': | |||
f = torch.tanh | |||
elif name == 'identity': | |||
f = lambda x: x | |||
elif name == 'sigmoid': | |||
f = torch.sigmoid | |||
return f | |||
@property | |||
def num_parameters(self): | |||
def size(p): | |||
return np.prod(p.size()) | |||
return sum([size(param) for param in self.parameters()]) | |||
def reset_parameters(self): | |||
init_range = 0.025 | |||
# init_range = 0.025 if self.args.mode == 'train' else 0.04 | |||
for param in self.parameters(): | |||
param.data.uniform_(-init_range, init_range) | |||
self.decoder.bias.data.fill_(0) | |||
def predict(self, word_seq): | |||
""" | |||
:param word_seq: torch.LongTensor, [batch_size, seq_len] | |||
:return predict: dict of torch.LongTensor, [batch_size, seq_len] | |||
""" | |||
output = self(word_seq) | |||
_, predict = output['pred'].max(dim=1) | |||
return {'pred': predict} |
@@ -0,0 +1,385 @@ | |||
# Code Modified from https://github.com/carpedm20/ENAS-pytorch | |||
import os | |||
import time | |||
from datetime import datetime | |||
from datetime import timedelta | |||
import numpy as np | |||
import torch | |||
import math | |||
from torch import nn | |||
try: | |||
from tqdm.autonotebook import tqdm | |||
except: | |||
from fastNLP.core.utils import pseudo_tqdm as tqdm | |||
from fastNLP.core.batch import Batch | |||
from fastNLP.core.callback import CallbackManager, CallbackException | |||
from fastNLP.core.dataset import DataSet | |||
from fastNLP.core.utils import CheckError | |||
from fastNLP.core.utils import _move_dict_value_to_device | |||
import fastNLP | |||
import fastNLP.models.enas_utils as utils | |||
from fastNLP.core.utils import _build_args | |||
from torch.optim import Adam | |||
def _get_no_grad_ctx_mgr(): | |||
"""Returns a the `torch.no_grad` context manager for PyTorch version >= | |||
0.4, or a no-op context manager otherwise. | |||
""" | |||
return torch.no_grad() | |||
class ENASTrainer(fastNLP.Trainer): | |||
"""A class to wrap training code.""" | |||
def __init__(self, train_data, model, controller, **kwargs): | |||
"""Constructor for training algorithm. | |||
:param DataSet train_data: the training data | |||
:param torch.nn.modules.module model: a PyTorch model | |||
:param torch.nn.modules.module controller: a PyTorch model | |||
""" | |||
self.final_epochs = kwargs['final_epochs'] | |||
kwargs.pop('final_epochs') | |||
super(ENASTrainer, self).__init__(train_data, model, **kwargs) | |||
self.controller_step = 0 | |||
self.shared_step = 0 | |||
self.max_length = 35 | |||
self.shared = model | |||
self.controller = controller | |||
self.shared_optim = Adam( | |||
self.shared.parameters(), | |||
lr=20.0, | |||
weight_decay=1e-7) | |||
self.controller_optim = Adam( | |||
self.controller.parameters(), | |||
lr=3.5e-4) | |||
def train(self, load_best_model=True): | |||
""" | |||
:param bool load_best_model: 该参数只有在初始化提供了dev_data的情况下有效,如果True, trainer将在返回之前重新加载dev表现 | |||
最好的模型参数。 | |||
:return results: 返回一个字典类型的数据, 内含以下内容:: | |||
seconds: float, 表示训练时长 | |||
以下三个内容只有在提供了dev_data的情况下会有。 | |||
best_eval: Dict of Dict, 表示evaluation的结果 | |||
best_epoch: int,在第几个epoch取得的最佳值 | |||
best_step: int, 在第几个step(batch)更新取得的最佳值 | |||
""" | |||
results = {} | |||
if self.n_epochs <= 0: | |||
print(f"training epoch is {self.n_epochs}, nothing was done.") | |||
results['seconds'] = 0. | |||
return results | |||
try: | |||
if torch.cuda.is_available() and self.use_cuda: | |||
self.model = self.model.cuda() | |||
self._model_device = self.model.parameters().__next__().device | |||
self._mode(self.model, is_test=False) | |||
self.start_time = str(datetime.now().strftime('%Y-%m-%d-%H-%M-%S')) | |||
start_time = time.time() | |||
print("training epochs started " + self.start_time, flush=True) | |||
try: | |||
self.callback_manager.on_train_begin() | |||
self._train() | |||
self.callback_manager.on_train_end(self.model) | |||
except (CallbackException, KeyboardInterrupt) as e: | |||
self.callback_manager.on_exception(e, self.model) | |||
if self.dev_data is not None: | |||
print("\nIn Epoch:{}/Step:{}, got best dev performance:".format(self.best_dev_epoch, self.best_dev_step) + | |||
self.tester._format_eval_results(self.best_dev_perf),) | |||
results['best_eval'] = self.best_dev_perf | |||
results['best_epoch'] = self.best_dev_epoch | |||
results['best_step'] = self.best_dev_step | |||
if load_best_model: | |||
model_name = "best_" + "_".join([self.model.__class__.__name__, self.metric_key, self.start_time]) | |||
load_succeed = self._load_model(self.model, model_name) | |||
if load_succeed: | |||
print("Reloaded the best model.") | |||
else: | |||
print("Fail to reload best model.") | |||
finally: | |||
pass | |||
results['seconds'] = round(time.time() - start_time, 2) | |||
return results | |||
def _train(self): | |||
if not self.use_tqdm: | |||
from fastNLP.core.utils import pseudo_tqdm as inner_tqdm | |||
else: | |||
inner_tqdm = tqdm | |||
self.step = 0 | |||
start = time.time() | |||
total_steps = (len(self.train_data) // self.batch_size + int( | |||
len(self.train_data) % self.batch_size != 0)) * self.n_epochs | |||
with inner_tqdm(total=total_steps, postfix='loss:{0:<6.5f}', leave=False, dynamic_ncols=True) as pbar: | |||
avg_loss = 0 | |||
data_iterator = Batch(self.train_data, batch_size=self.batch_size, sampler=self.sampler, as_numpy=False, | |||
prefetch=self.prefetch) | |||
for epoch in range(1, self.n_epochs+1): | |||
pbar.set_description_str(desc="Epoch {}/{}".format(epoch, self.n_epochs)) | |||
last_stage = (epoch > self.n_epochs + 1 - self.final_epochs) | |||
if epoch == self.n_epochs + 1 - self.final_epochs: | |||
print('Entering the final stage. (Only train the selected structure)') | |||
# early stopping | |||
self.callback_manager.on_epoch_begin(epoch, self.n_epochs) | |||
# 1. Training the shared parameters omega of the child models | |||
self.train_shared(pbar) | |||
# 2. Training the controller parameters theta | |||
if not last_stage: | |||
self.train_controller() | |||
if ((self.validate_every > 0 and self.step % self.validate_every == 0) or | |||
(self.validate_every < 0 and self.step % len(data_iterator) == 0)) \ | |||
and self.dev_data is not None: | |||
if not last_stage: | |||
self.derive() | |||
eval_res = self._do_validation(epoch=epoch, step=self.step) | |||
eval_str = "Evaluation at Epoch {}/{}. Step:{}/{}. ".format(epoch, self.n_epochs, self.step, | |||
total_steps) + \ | |||
self.tester._format_eval_results(eval_res) | |||
pbar.write(eval_str) | |||
# lr decay; early stopping | |||
self.callback_manager.on_epoch_end(epoch, self.n_epochs, self.optimizer) | |||
# =============== epochs end =================== # | |||
pbar.close() | |||
# ============ tqdm end ============== # | |||
def get_loss(self, inputs, targets, hidden, dags): | |||
"""Computes the loss for the same batch for M models. | |||
This amounts to an estimate of the loss, which is turned into an | |||
estimate for the gradients of the shared model. | |||
""" | |||
if not isinstance(dags, list): | |||
dags = [dags] | |||
loss = 0 | |||
for dag in dags: | |||
self.shared.setDAG(dag) | |||
inputs = _build_args(self.shared.forward, **inputs) | |||
inputs['hidden'] = hidden | |||
result = self.shared(**inputs) | |||
output, hidden, extra_out = result['pred'], result['hidden'], result['extra_out'] | |||
self.callback_manager.on_loss_begin(targets, result) | |||
sample_loss = self._compute_loss(result, targets) | |||
loss += sample_loss | |||
assert len(dags) == 1, 'there are multiple `hidden` for multple `dags`' | |||
return loss, hidden, extra_out | |||
def train_shared(self, pbar=None, max_step=None, dag=None): | |||
"""Train the language model for 400 steps of minibatches of 64 | |||
examples. | |||
Args: | |||
max_step: Used to run extra training steps as a warm-up. | |||
dag: If not None, is used instead of calling sample(). | |||
BPTT is truncated at 35 timesteps. | |||
For each weight update, gradients are estimated by sampling M models | |||
from the fixed controller policy, and averaging their gradients | |||
computed on a batch of training data. | |||
""" | |||
model = self.shared | |||
model.train() | |||
self.controller.eval() | |||
hidden = self.shared.init_hidden(self.batch_size) | |||
abs_max_grad = 0 | |||
abs_max_hidden_norm = 0 | |||
step = 0 | |||
raw_total_loss = 0 | |||
total_loss = 0 | |||
train_idx = 0 | |||
avg_loss = 0 | |||
data_iterator = Batch(self.train_data, batch_size=self.batch_size, sampler=self.sampler, as_numpy=False, | |||
prefetch=self.prefetch) | |||
for batch_x, batch_y in data_iterator: | |||
_move_dict_value_to_device(batch_x, batch_y, device=self._model_device) | |||
indices = data_iterator.get_batch_indices() | |||
# negative sampling; replace unknown; re-weight batch_y | |||
self.callback_manager.on_batch_begin(batch_x, batch_y, indices) | |||
# prediction = self._data_forward(self.model, batch_x) | |||
dags = self.controller.sample(1) | |||
inputs, targets = batch_x, batch_y | |||
# self.callback_manager.on_loss_begin(batch_y, prediction) | |||
loss, hidden, extra_out = self.get_loss(inputs, | |||
targets, | |||
hidden, | |||
dags) | |||
hidden.detach_() | |||
avg_loss += loss.item() | |||
# Is loss NaN or inf? requires_grad = False | |||
self.callback_manager.on_backward_begin(loss, self.model) | |||
self._grad_backward(loss) | |||
self.callback_manager.on_backward_end(self.model) | |||
self._update() | |||
self.callback_manager.on_step_end(self.optimizer) | |||
if (self.step+1) % self.print_every == 0: | |||
if self.use_tqdm: | |||
print_output = "loss:{0:<6.5f}".format(avg_loss / self.print_every) | |||
pbar.update(self.print_every) | |||
else: | |||
end = time.time() | |||
diff = timedelta(seconds=round(end - start)) | |||
print_output = "[epoch: {:>3} step: {:>4}] train loss: {:>4.6} time: {}".format( | |||
epoch, self.step, avg_loss, diff) | |||
pbar.set_postfix_str(print_output) | |||
avg_loss = 0 | |||
self.step += 1 | |||
step += 1 | |||
self.shared_step += 1 | |||
self.callback_manager.on_batch_end() | |||
# ================= mini-batch end ==================== # | |||
def get_reward(self, dag, entropies, hidden, valid_idx=0): | |||
"""Computes the perplexity of a single sampled model on a minibatch of | |||
validation data. | |||
""" | |||
if not isinstance(entropies, np.ndarray): | |||
entropies = entropies.data.cpu().numpy() | |||
data_iterator = Batch(self.dev_data, batch_size=self.batch_size, sampler=self.sampler, as_numpy=False, | |||
prefetch=self.prefetch) | |||
for inputs, targets in data_iterator: | |||
valid_loss, hidden, _ = self.get_loss(inputs, targets, hidden, dag) | |||
valid_loss = utils.to_item(valid_loss.data) | |||
valid_ppl = math.exp(valid_loss) | |||
R = 80 / valid_ppl | |||
rewards = R + 1e-4 * entropies | |||
return rewards, hidden | |||
def train_controller(self): | |||
"""Fixes the shared parameters and updates the controller parameters. | |||
The controller is updated with a score function gradient estimator | |||
(i.e., REINFORCE), with the reward being c/valid_ppl, where valid_ppl | |||
is computed on a minibatch of validation data. | |||
A moving average baseline is used. | |||
The controller is trained for 2000 steps per epoch (i.e., | |||
first (Train Shared) phase -> second (Train Controller) phase). | |||
""" | |||
model = self.controller | |||
model.train() | |||
# Why can't we call shared.eval() here? Leads to loss | |||
# being uniformly zero for the controller. | |||
# self.shared.eval() | |||
avg_reward_base = None | |||
baseline = None | |||
adv_history = [] | |||
entropy_history = [] | |||
reward_history = [] | |||
hidden = self.shared.init_hidden(self.batch_size) | |||
total_loss = 0 | |||
valid_idx = 0 | |||
for step in range(20): | |||
# sample models | |||
dags, log_probs, entropies = self.controller.sample( | |||
with_details=True) | |||
# calculate reward | |||
np_entropies = entropies.data.cpu().numpy() | |||
# No gradients should be backpropagated to the | |||
# shared model during controller training, obviously. | |||
with _get_no_grad_ctx_mgr(): | |||
rewards, hidden = self.get_reward(dags, | |||
np_entropies, | |||
hidden, | |||
valid_idx) | |||
reward_history.extend(rewards) | |||
entropy_history.extend(np_entropies) | |||
# moving average baseline | |||
if baseline is None: | |||
baseline = rewards | |||
else: | |||
decay = 0.95 | |||
baseline = decay * baseline + (1 - decay) * rewards | |||
adv = rewards - baseline | |||
adv_history.extend(adv) | |||
# policy loss | |||
loss = -log_probs*utils.get_variable(adv, | |||
self.use_cuda, | |||
requires_grad=False) | |||
loss = loss.sum() # or loss.mean() | |||
# update | |||
self.controller_optim.zero_grad() | |||
loss.backward() | |||
self.controller_optim.step() | |||
total_loss += utils.to_item(loss.data) | |||
if ((step % 50) == 0) and (step > 0): | |||
reward_history, adv_history, entropy_history = [], [], [] | |||
total_loss = 0 | |||
self.controller_step += 1 | |||
# prev_valid_idx = valid_idx | |||
# valid_idx = ((valid_idx + self.max_length) % | |||
# (self.valid_data.size(0) - 1)) | |||
# # Whenever we wrap around to the beginning of the | |||
# # validation data, we reset the hidden states. | |||
# if prev_valid_idx > valid_idx: | |||
# hidden = self.shared.init_hidden(self.batch_size) | |||
def derive(self, sample_num=10, valid_idx=0): | |||
"""We are always deriving based on the very first batch | |||
of validation data? This seems wrong... | |||
""" | |||
hidden = self.shared.init_hidden(self.batch_size) | |||
dags, _, entropies = self.controller.sample(sample_num, | |||
with_details=True) | |||
max_R = 0 | |||
best_dag = None | |||
for dag in dags: | |||
R, _ = self.get_reward(dag, entropies, hidden, valid_idx) | |||
if R.max() > max_R: | |||
max_R = R.max() | |||
best_dag = dag | |||
self.model.setDAG(best_dag) |
@@ -0,0 +1,56 @@ | |||
# Code Modified from https://github.com/carpedm20/ENAS-pytorch | |||
from __future__ import print_function | |||
from collections import defaultdict | |||
import collections | |||
from datetime import datetime | |||
import os | |||
import json | |||
import numpy as np | |||
import torch | |||
from torch.autograd import Variable | |||
def detach(h): | |||
if type(h) == Variable: | |||
return Variable(h.data) | |||
else: | |||
return tuple(detach(v) for v in h) | |||
def get_variable(inputs, cuda=False, **kwargs): | |||
if type(inputs) in [list, np.ndarray]: | |||
inputs = torch.Tensor(inputs) | |||
if cuda: | |||
out = Variable(inputs.cuda(), **kwargs) | |||
else: | |||
out = Variable(inputs, **kwargs) | |||
return out | |||
def update_lr(optimizer, lr): | |||
for param_group in optimizer.param_groups: | |||
param_group['lr'] = lr | |||
Node = collections.namedtuple('Node', ['id', 'name']) | |||
class keydefaultdict(defaultdict): | |||
def __missing__(self, key): | |||
if self.default_factory is None: | |||
raise KeyError(key) | |||
else: | |||
ret = self[key] = self.default_factory(key) | |||
return ret | |||
def to_item(x): | |||
"""Converts x, possibly scalar and possibly tensor, to a Python scalar.""" | |||
if isinstance(x, (float, int)): | |||
return x | |||
if float(torch.__version__[0:3]) < 0.4: | |||
assert (x.dim() == 1) and (len(x) == 1) | |||
return x[0] | |||
return x.item() |
@@ -0,0 +1,112 @@ | |||
import unittest | |||
from fastNLP import DataSet | |||
from fastNLP import Instance | |||
from fastNLP import Vocabulary | |||
from fastNLP.core.losses import CrossEntropyLoss | |||
from fastNLP.core.metrics import AccuracyMetric | |||
class TestENAS(unittest.TestCase): | |||
def testENAS(self): | |||
# 从csv读取数据到DataSet | |||
sample_path = "tutorials/sample_data/tutorial_sample_dataset.csv" | |||
dataset = DataSet.read_csv(sample_path, headers=('raw_sentence', 'label'), | |||
sep='\t') | |||
print(len(dataset)) | |||
print(dataset[0]) | |||
print(dataset[-3]) | |||
dataset.append(Instance(raw_sentence='fake data', label='0')) | |||
# 将所有数字转为小写 | |||
dataset.apply(lambda x: x['raw_sentence'].lower(), new_field_name='raw_sentence') | |||
# label转int | |||
dataset.apply(lambda x: int(x['label']), new_field_name='label') | |||
# 使用空格分割句子 | |||
def split_sent(ins): | |||
return ins['raw_sentence'].split() | |||
dataset.apply(split_sent, new_field_name='words') | |||
# 增加长度信息 | |||
dataset.apply(lambda x: len(x['words']), new_field_name='seq_len') | |||
print(len(dataset)) | |||
print(dataset[0]) | |||
# DataSet.drop(func)筛除数据 | |||
dataset.drop(lambda x: x['seq_len'] <= 3) | |||
print(len(dataset)) | |||
# 设置DataSet中,哪些field要转为tensor | |||
# set target,loss或evaluate中的golden,计算loss,模型评估时使用 | |||
dataset.set_target("label") | |||
# set input,模型forward时使用 | |||
dataset.set_input("words", "seq_len") | |||
# 分出测试集、训练集 | |||
test_data, train_data = dataset.split(0.5) | |||
print(len(test_data)) | |||
print(len(train_data)) | |||
# 构建词表, Vocabulary.add(word) | |||
vocab = Vocabulary(min_freq=2) | |||
train_data.apply(lambda x: [vocab.add(word) for word in x['words']]) | |||
vocab.build_vocab() | |||
# index句子, Vocabulary.to_index(word) | |||
train_data.apply(lambda x: [vocab.to_index(word) for word in x['words']], new_field_name='words') | |||
test_data.apply(lambda x: [vocab.to_index(word) for word in x['words']], new_field_name='words') | |||
print(test_data[0]) | |||
# 如果你们需要做强化学习或者GAN之类的项目,你们也可以使用这些数据预处理的工具 | |||
from fastNLP.core.batch import Batch | |||
from fastNLP.core.sampler import RandomSampler | |||
batch_iterator = Batch(dataset=train_data, batch_size=2, sampler=RandomSampler()) | |||
for batch_x, batch_y in batch_iterator: | |||
print("batch_x has: ", batch_x) | |||
print("batch_y has: ", batch_y) | |||
break | |||
from fastNLP.models.enas_model import ENASModel | |||
from fastNLP.models.enas_controller import Controller | |||
model = ENASModel(embed_num=len(vocab), num_classes=5) | |||
controller = Controller() | |||
from fastNLP.models.enas_trainer import ENASTrainer | |||
from copy import deepcopy | |||
# 更改DataSet中对应field的名称,要以模型的forward等参数名一致 | |||
train_data.rename_field('words', 'word_seq') # input field 与 forward 参数一致 | |||
train_data.rename_field('label', 'label_seq') | |||
test_data.rename_field('words', 'word_seq') | |||
test_data.rename_field('label', 'label_seq') | |||
loss = CrossEntropyLoss(pred="output", target="label_seq") | |||
metric = AccuracyMetric(pred="predict", target="label_seq") | |||
trainer = ENASTrainer(model=model, controller=controller, train_data=train_data, dev_data=test_data, | |||
loss=CrossEntropyLoss(pred="output", target="label_seq"), | |||
metrics=AccuracyMetric(pred="predict", target="label_seq"), | |||
check_code_level=-1, | |||
save_path=None, | |||
batch_size=32, | |||
print_every=1, | |||
n_epochs=3, | |||
final_epochs=1) | |||
trainer.train() | |||
print('Train finished!') | |||
# 调用Tester在test_data上评价效果 | |||
from fastNLP import Tester | |||
tester = Tester(data=test_data, model=model, metrics=AccuracyMetric(pred="predict", target="label_seq"), | |||
batch_size=4) | |||
acc = tester.test() | |||
print(acc) | |||
if __name__ == '__main__': | |||
unittest.main() |