简单讲讲 mobile aloha 训练阶段的代码-CVAE和Transformer神经网络 (一)
今天简单聊聊 mobile aloha CVAE 和 Transformer 神经网络的构建;
代码主要在 detr_vae.py 里的 build 函数(detr表示DEtection TRansformer);
build 函数主要创建下面的几个神经网络:
backbones 神经网络
用来处理机械臂摄像头数据, 在mujoco虚拟环境中有3个摄像头, 分别是 top(上面的摄像头), left_wrist(左机械臂手腕的摄像头), right_wrist(右机械臂手腕摄像头);
backbones = [] for _ in args.camera_names: backbone = build_backbone(args) backbones.append(backbone)
在训练的时候 backbones 神经网络会把摄像头数据转为 features(特征值)和pos(位置嵌入), features的形状是(8, 512, 15, 20)的张量, pos的形状是(1, 512, 15, 60)的张量, 后续会进行展平和变换, 转到512维空间中, 变成 (sequence, bs, 512) 的形状, sequence是输入序列, bs是批次, 512是维度;
backbones的细节我们今天先跳过;
transformer神经网络
下面的代码初始化 Transformer 模型的编码器和解码器, 以及用于归一化的层;
encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, activation, normalize_before) encoder_norm = nn.LayerNorm(d_model) if normalize_before else None self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm) decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout, activation, normalize_before) decoder_norm = nn.LayerNorm(d_model) self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm, return_intermediate=return_intermediate_dec)
之后用Xavier 均匀分布对神经网络的权重进行初始化;
for p in self.parameters(): if p.dim() > 1: nn.init.xavier_uniform_(p)
TransformerEncoderLayer的代码:
self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) # Implementation of Feedforward model self.linear1 = nn.Linear(d_model, dim_feedforward) self.dropout = nn.Dropout(dropout) self.linear2 = nn.Linear(dim_feedforward, d_model) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) self.activation = _get_activation_fn(activation) self.normalize_before = normalize_before
使用 pytorch 的 MultiheadAttention 构建注意力层 ;
使用 pytorch 的 Linear, Dropout, LayerNorm 构建前馈神经网络;
设置激活函数;
设置一个标志, 是否在输入之前进行归一化;
下面我们会介绍 transformer enconde layer 怎么调用这些层;
在前向传播中, Transformer接受的参数是摄像头数据, CVAE隐变量, 机械臂自身的感知数据:
hs = self.transformer(src, None, self.query_embed.weight, pos, latent_input, proprio_input, self.additional_pos_embed.weight)[0]
返回值hs(hiden states)是Transformer神经网络decoder之后的值, 形状是(8, 100, 512)的张量, 之后, 这个张量经过线性变换, 从512维到16维;
a_hat = self.action_head(hs)
a_hat的形状是(8, 100, 16), 8是批次, 100是帧数, 16是机械臂各个部件的动作, 每个机械臂有8个关节, 两个机械臂加起来是16;
我们看看transformer的逻辑:
transformer在调用transformer encoder进行前向传播之前先把输入数据进行展平, 变换, 拼接, 转换到512维;
# flatten NxCxHxW to HWxNxC bs, c, h, w = src.shape src = src.flatten(2).permute(2, 0, 1) pos_embed = pos_embed.flatten(2).permute(2, 0, 1).repeat(1, bs, 1) query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1) # mask = mask.flatten(1) additional_pos_embed = additional_pos_embed.unsqueeze(1).repeat(1, bs, 1) # seq, bs, dim pos_embed = torch.cat([additional_pos_embed, pos_embed], axis=0) addition_input = torch.stack([latent_input, proprio_input], axis=0) src = torch.cat([addition_input, src], axis=0)
这里得到的src的是形状(902, 8, 512)的张量, 包含了3个摄像头的数据(300*3, 8, 512) 加 隐变量(1, 8, 512) 加 机械臂自身感知(1, 8, 512), 这个张量会被传到transformer的encoder中进行计算;
memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed) hs = self.decoder(tgt, memory, memory_key_padding_mask=mask, pos=pos_embed, query_pos=query_embed)
得到的 memory 是记忆张量; 我们先看看 transformer encoder 的实现:
def forward(self, src, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None): output = src for layer in self.layers: output = layer(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask, pos=pos) if self.norm is not None: output = self.norm(output) return output
在前向传播中, 会调用每个 transformer encoder layer 进行前向传播, 然后判断是否做归一化, 返回输出结果; transfoemer encoder包含了4层 transformer encoder layer;
我们看看 transformer encoder layer 的实现:
q = k = self.with_pos_embed(src, pos) src2 = self.self_attn(q, k, value=src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0] src = src + self.dropout1(src2) src = self.norm1(src) src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) src = src + self.dropout2(src2) src = self.norm2(src) return src
先生成query和key, 形状是(902, 8, 512)的张量, 512是维度;
q = k = self.with_pos_embed(src, pos)
调用 self.self_attn 来计算注意力, (参考前面的代码, self.self_attn是nn.MultiheadAttention(512, 8, dropout));
src2 = self.self_attn(q, k, value=src, attn_mask=src_mask,
调用 self.dropout1 随机丢弃一些神经输出, 然后和原始 src 相加, 实现残差连接(self.dropout1是nn.Dropout(dropout));
src = src + self.dropout1(src2)
之后, 调用 self.norm1 进行归一化(self.norm1是nn.LayerNorm(512));
src = self.norm1(src)
对 src 调用一次线性层(self.linear1), 应用激活函数(self.activation), 应用dropout(self.dropout), 传播到第二个线性层(self.linear2), 其中: self.linear1 = nn.Linear(512, 512), self.activation=RELU, self.dropout = nn.Dropout(dropout), self.linear2 = nn.Linear(512, 512);
src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
再进行一次残差连接:
src = src + self.dropout2(src2)
再进行一次归一化:
src = self.norm2(src)
返回结果;
我们回到前面的代码,
memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed)
transformer encoder 之后, 我们得到了记忆张量 memory, 形状是 (902, 8, 512);
记忆张量经过 transformer decoder, 得到 hs(hiden states)
hs = self.decoder(tgt, memory, memory_key_padding_mask=mask,
pos=pos_embed, query_pos=query_embed)
hs 再经过一个线性变换, 得到预测的机械臂的action
a_hat = self.action_head(hs)
hs是(8, 100, 512), a_hat是(8, 100, 16), 16表示机械臂的动作;
总体来说, 通过 transformer 的 encoder 和 decoder 实现了 CVAE 的解码器;
今天暂时到这里, 下次我们接着看 transformer 的 decoder 和 CVAE encoder;
浙公网安备 33010602011771号