Spaces:
Runtime error
Runtime error
File size: 5,206 Bytes
bc32eea |
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 116 |
import torch
from torch import nn
from .transformer import BasicTransformerModel
from models import BaseModel
from .util.generation import autoregressive_generation_multimodal
class TransformerModel(BaseModel):
def __init__(self, opt):
super().__init__(opt)
opt=self.opt
input_mods = self.input_mods
output_mods = self.output_mods
dins = self.dins
douts = self.douts
input_lengths = self.input_lengths
self.input_mod_nets = []
self.output_mod_nets = []
self.module_names = []
for i, mod in enumerate(input_mods):
net = BasicTransformerModel(opt.dhid, dins[i], opt.nhead, opt.dhid, 2, opt.dropout, self.device, use_pos_emb=True, input_length=input_lengths[i], use_x_transformers=opt.use_x_transformers, opt=opt)
name = "_input_"+mod
setattr(self,"net"+name, net)
self.input_mod_nets.append(net)
self.module_names.append(name)
for i, mod in enumerate(output_mods):
net = BasicTransformerModel(douts[i], opt.dhid, opt.nhead, opt.dhid, opt.nlayers, opt.dropout, self.device, use_pos_emb=opt.use_pos_emb_output, input_length=sum(input_lengths), use_x_transformers=opt.use_x_transformers, opt=opt)
# net = BasicTransformerModel(douts[i], opt.dhid, opt.nhead, opt.dhid, opt.nlayers, opt.dropout, self.device, use_pos_emb=True, input_length=sum(input_lengths))
name = "_output_"+mod
setattr(self,"net"+name, net)
self.output_mod_nets.append(net)
self.module_names.append(name)
#This is feature creep. Will remove soon
# if self.opt.generate_attention_masks:
self.generate_full_masks()
self.inputs = []
self.targets = []
self.criterion = nn.MSELoss()
def name(self):
return "Transformer"
@staticmethod
def modify_commandline_options(parser, opt):
parser.add_argument('--dhid', type=int, default=512)
parser.add_argument('--nlayers', type=int, default=6)
parser.add_argument('--nhead', type=int, default=8)
parser.add_argument('--dropout', type=float, default=0.1)
parser.add_argument('--use_pos_emb_output', action='store_true', help="whether to use positional embeddings for output modality transformers")
parser.add_argument('--use_rotary_pos_emb', action='store_true', help="whether to use rotary position embeddings")
parser.add_argument('--use_x_transformers', action='store_true', help="whether to use rotary position embeddings")
# parser.add_argument('--generate_attention_masks', action='store_true', help="whether to generate the masks (but right now they are full masks, so it's not necessary")
return parser
def generate_full_masks(self):
input_mods = self.input_mods
output_mods = self.output_mods
input_lengths = self.input_lengths
self.src_masks = []
for i, mod in enumerate(input_mods):
mask = torch.zeros(input_lengths[i],input_lengths[i])
self.register_buffer('src_mask_'+str(i), mask)
self.src_masks.append(mask)
self.output_masks = []
for i, mod in enumerate(output_mods):
mask = torch.zeros(sum(input_lengths),sum(input_lengths))
self.register_buffer('out_mask_'+str(i), mask)
self.output_masks.append(mask)
def forward(self, data):
# in lightning, forward defines the prediction/inference actions
latents = []
for i, mod in enumerate(self.input_mods):
latents.append(self.input_mod_nets[i].forward(data[i]))
latent = torch.cat(latents)
outputs = []
for i, mod in enumerate(self.output_mods):
output = self.output_mod_nets[i].forward(latent)[:self.output_lengths[i]]
outputs.append(output)
#import pdb;pdb.set_trace()
return outputs
def training_step(self, batch, batch_idx):
self.set_inputs(batch)
#print(self.inputs)
latents = []
for i, mod in enumerate(self.input_mods):
latents.append(self.input_mod_nets[i].forward(self.inputs[i]))
latent = torch.cat(latents)
loss_mse = 0
for i, mod in enumerate(self.output_mods):
output = self.output_mod_nets[i].forward(latent)[:self.output_lengths[i]]
#print(output)
loss_mse += self.criterion(output, self.targets[i])
#loss_mse += self.criterion(output, self.targets[i]).detach()
#print(loss_mse)
#if self.opt.precision == 16:
# loss_mse *= 100 # loss scaling
self.log('mse_loss', loss_mse)
return loss_mse
#return torch.tensor(0.0, dtype=torch.float32, requires_grad=True)
#def configure_optimizers(self):
# print("HIIIIIIIIIIIIIIIIII")
# optimizer = torch.optim.Adam(self.parameters(), lr=self.opt.learning_rate)
# return [optimizer]
#def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx,
# optimizer_closure, on_tpu, using_native_amp, using_lbfgs):
# optimizer.zero_grad()
|