import torch, torch.nn as nn
class A(nn.Module):
def __init__(self):
super().__init__()
self.embed_src = nn.Embedding(2, 10)
self.embed_target = nn.Embedding(2, 10)
self.transformer = nn.Transformer(10, 2)
self.lin = nn.Linear(10, 2)
self.softmax = nn.Softmax(dim=-1)
def forward(self, inp, tgt):
embed_src = self.embed_src(inp)
embed_target = self.embed_target(tgt)
# print(embedding.shape)
output = self.transformer(embed_src.view(len(inp), 1, -1), embed_target.view((len(tgt)), 1, -1))
output = self.lin(output)
print('output.shape', output.shape, 'output', output)
print('embed_target', embed_target.view(2, 1, -1))
print('softmax probabilities', self.softmax(output))
return output.permute(0, 2, 1)
model = A()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
dataset = [
[([0, 1, 0, 1, 0, 1]), ([[0], [1]])],
[([1, 0, 1, 0, 1, 0]), ([[1], [0]])],
[([0, 0, 0, 0, 0, 0]), ([[0], [0]])],
[([1, 1, 1, 1, 1, 1]), ([[1], [1]])],
[([1, 1, 0, 0, 1, 1]), ([[0], [0]])],
[([0, 0, 1, 1, 0, 0]), ([[1], [1]])]
]
tensor_dataset = []
for i in range(6):
tensor_dataset.append([torch.tensor(dataset[i][0]), torch.tensor(dataset[i][1])])
tensor_dataset
[
[tensor([0, 1, 0, 1, 0, 1]), tensor([[0], [1]])],
[tensor([1, 0, 1, 0, 1, 0]), tensor([[1], [0]])],
[tensor([0, 0, 0, 0, 0, 0]), tensor([[0], [0]])],
[tensor([1, 1, 1, 1, 1, 1]), tensor([[1], [1]])],
[tensor([1, 1, 0, 0, 1, 1]), tensor([[0], [0]])],
[tensor([0, 0, 1, 1, 0, 0]), tensor([[1], [1]])]
]
criterion = nn.CrossEntropyLoss()
for i in range(1000):
optimizer.zero_grad()
print(tensor_dataset[i%6][0], tensor_dataset[i%6][1])
loss = criterion(model(tensor_dataset[i%6][0], tensor_dataset[i%6][1]), tensor_dataset[i%6][1])
print(loss)
loss.backward()
optimizer.step()
do I need to set self.embed_target requires_grad as False?
is output.permute(0, 2, 1) correct way to use with nn.CrossEntropyLoss?