I am trying to run my model on the GPU.
Here is the what I have so far:
from torch.utils.data import DataLoader
import torch.optim as optim
import numpy as np
from torch import nn
import torch
from packages.ann import ANN
from packages.cnn import CNN
class MODEL(nn.Module):
def __init__(self, input_dim, input_length, output_dim, kernel_size=16):
super(self, MODEL).__init__()
self.ann = ANN(input_dim, output_dim ,kernel_size, dropout = 0.2)
self.cnn = CNN(input_dim, output_dim ,kernel_size, dropout = 0.2)
def forward(self, x):
ann_output = self.ann(x)[:, -1]
cnn_output = self.cnn(x)[:, -1]
return torch.cat((ann_output, cnn_output), dim=1)
# set env variable for gpus
os.environ["CUDA_VISIBLE_DEVICES"]="0"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Call the model
model = MODEL(input_dim, input_length, output_dim, kernel_size)
model = model.to(device, dtype=torch.float)
optimizer = optim.Adam(model.parameters(), lr=0.005, weight_decay=0.85)
triplet_loss = nn.TripletMarginLoss(margin=1)
model.train()
#`train_data` are of type numpy
loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
epochs = 10
for epoch in range():
running_loss = []
for step, (anchor, positive, negative) in enumerate(loader):
anchor = anchor.to(device)
positive = positive.to(device)
negative = negative.to(device)
print(type(anchor)) # <class 'torch.Tensor'>
optimizer.zero_grad()
anchor_output = model(anchor) # The error occurs here.
positive_output = model(positive)
negative_output = model(negative)
loss = triplet_loss(anchor_output, positive_output, negative_output)
loss.backward()
optimizer.step()
running_loss.append(loss.cpu().detach().numpy())
I get the following error in the second for loop, on line anchor_output = model(anchor)
TypeError: Cannot convert torch.dtype to numpy.dtype