Spaces:
Runtime error
Runtime error
File size: 4,552 Bytes
8fc2b4e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
import numpy as np
import torch
import torch.nn.functional as F
from cliport.models.core.attention import Attention
import cliport.models as models
import cliport.models.core.fusion as fusion
class TwoStreamAttentionLangFusion(Attention):
"""Two Stream Language-Conditioned Attention (a.k.a Pick) module."""
def __init__(self, stream_fcn, in_shape, n_rotations, preprocess, cfg, device):
self.fusion_type = cfg['train']['attn_stream_fusion_type']
super().__init__(stream_fcn, in_shape, n_rotations, preprocess, cfg, device)
def _build_nets(self):
stream_one_fcn, stream_two_fcn = self.stream_fcn
stream_one_model = models.names[stream_one_fcn] # resnet_lat.REsNet45_10s
stream_two_model = models.names[stream_two_fcn] # clip_ligunet_lat.CLIP_LIGUnet_lat
self.attn_stream_one = stream_one_model(self.in_shape, 1, self.cfg, self.device, self.preprocess)
self.attn_stream_two = stream_two_model(self.in_shape, 1, self.cfg, self.device, self.preprocess)
self.fusion = fusion.names[self.fusion_type](input_dim=1)
print(f"Attn FCN - Stream One: {stream_one_fcn}, Stream Two: {stream_two_fcn}, Stream Fusion: {self.fusion_type}")
def attend(self, x, l):
x1 = self.attn_stream_one(x)
x2 = self.attn_stream_two(x, l)
x = self.fusion(x1, x2)
return x
def forward(self, inp_img, lang_goal, softmax=True):
"""Forward pass."""
if len(inp_img.shape) < 4:
inp_img = inp_img[None]
if type(inp_img) is not torch.Tensor:
in_data = inp_img # .reshape(in_shape)
in_tens = torch.from_numpy(in_data.copy()).to(dtype=torch.float, device=self.device) # [B W H 6]
else:
in_data = inp_img
in_tens = in_data
# [B W H 6]
in_tens = torch.nn.functional.pad(in_tens, tuple(self.padding[[2,1,0]].reshape(-1)), mode='constant')
# Rotation pivot.
pv = np.array(in_tens.shape[1:3]) // 2
# Rotate input.
in_tens = in_tens.permute(0, 3, 1, 2) # [B 6 W H]
# in_tens = in_tens.repeat(self.n_rotations, 1, 1, 1)
# make n copies, but keep batchsize
in_tens = [in_tens] * self.n_rotations
in_tens = self.rotator(in_tens, pivot=pv)
# Forward pass.
logits = self.attend(torch.cat(in_tens, dim=0), lang_goal)
# Rotate back output.
logits = self.rotator([logits], reverse=True, pivot=pv)
logits = torch.cat(logits, dim=0)
c0 = self.padding[:2, 0]
c1 = c0 + inp_img[0].shape[:2]
logits = logits[:, :, c0[0]:c1[0], c0[1]:c1[1]]
output_shape = logits.shape
# logits = logits.permute(1, 2, 3, 0) # [B W H 1]
output = logits.reshape(len(logits), -1)
if softmax:
output = F.softmax(output, dim=-1)
return output.view(output_shape)
class TwoStreamAttentionLangFusionLat(TwoStreamAttentionLangFusion):
"""Language-Conditioned Attention (a.k.a Pick) module with lateral connections."""
def __init__(self, stream_fcn, in_shape, n_rotations, preprocess, cfg, device):
self.fusion_type = cfg['train']['attn_stream_fusion_type']
super().__init__(stream_fcn, in_shape, n_rotations, preprocess, cfg, device)
def attend(self, x, l):
x1, lat = self.attn_stream_one(x)
x2 = self.attn_stream_two(x, lat, l)
x = self.fusion(x1, x2)
return x
class TwoStreamAttentionLangFusionLatReduce(TwoStreamAttentionLangFusion):
"""Language-Conditioned Attention (a.k.a Pick) module with lateral connections."""
def __init__(self, stream_fcn, in_shape, n_rotations, preprocess, cfg, device):
self.fusion_type = cfg['train']['attn_stream_fusion_type']
super().__init__(stream_fcn, in_shape, n_rotations, preprocess, cfg, device)
del self.attn_stream_one
del self.attn_stream_two
stream_one_fcn = 'plain_resnet_reduce_lat'
stream_one_model = models.names[stream_one_fcn]
stream_two_fcn = 'clip_ling'
stream_two_model = models.names[stream_two_fcn]
self.attn_stream_one = stream_one_model(self.in_shape, 1, self.cfg, self.device, self.preprocess)
self.attn_stream_two = stream_two_model(self.in_shape, 1, self.cfg, self.device, self.preprocess)
def attend(self, x, l):
x1, lat = self.attn_stream_one(x)
x2 = self.attn_stream_two(x, lat, l)
x = self.fusion(x1, x2)
return x |