From 9b411f00f3a3949a1ea96e4d5cf20710c44b2e74 Mon Sep 17 00:00:00 2001 From: Evan Arlian Date: Wed, 19 Oct 2022 14:06:46 +0700 Subject: [PATCH] remove data type conversion --- model2.py | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/model2.py b/model2.py index 55122e07f..74cce4270 100644 --- a/model2.py +++ b/model2.py @@ -2,7 +2,6 @@ from dataclasses import dataclass from typing import Optional, Generator - import numpy as np import torch from torch import nn @@ -10,9 +9,6 @@ from torch import Tensor -# TODO remove float(), another data types call to(x.dtype) - - @dataclass class ModelDimensions: n_mels: int @@ -62,7 +58,7 @@ def qkv_attention(self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor]) qk = q @ k if mask is not None: qk = qk + mask[:n_ctx, :n_ctx] - w = F.softmax(qk.float(), dim=-1).to(q.dtype) + w = F.softmax(qk, dim=-1) return (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2) @@ -105,7 +101,7 @@ def qkv_attention(self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor]) qk = q @ k if mask is not None: qk = qk + mask[:n_ctx, :n_ctx] - w = F.softmax(qk.float(), dim=-1).to(q.dtype) + w = F.softmax(qk, dim=-1) return (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2) @@ -178,7 +174,7 @@ def forward(self, x: Tensor): x = x.permute(0, 2, 1) assert x[0].size() == self.positional_embedding.size(), "incorrect audio shape" - x = (x + self.positional_embedding).to(x.dtype) + x = x + self.positional_embedding for block in self.blocks: x = block(x) @@ -225,15 +221,12 @@ def forward(self, x: Tensor, xa: Tensor, kv_cache: dict[int, Tensor]): self.token_embedding(x) + self.positional_embedding[offset : offset + x.size(-1)] ) - x = x.to(xa.dtype) for block in self.blocks: x = block(x, xa, mask=self.mask, kv_cache=kv_cache) x = self.ln(x) - logits = ( - x @ torch.transpose(self.token_embedding.weight.to(x.dtype), 0, 1) - ).float() + logits = x @ torch.transpose(self.token_embedding.weight, 0, 1) return logits