import matplotlib.image as mpimg
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import torch.utils.data as utils
from torchvision import transforms
import matplotlib.pyplot as plt
warnings.filterwarnings("ignore")
train_dir = data_dir + '/train/train/'
test_dir = data_dir + '/test/test/'
labels = pd.read_csv("../input/train.csv")
class ImageData(Dataset):
def __init__(self, df, data_dir, transform):
self.transform = transform
def __getitem__(self, index):
img_name = self.df.id[index]
label = self.df.has_cactus[index]
img_path = os.path.join(self.data_dir, img_name)
image = mpimg.imread(img_path)
image = self.transform(image)
data_transf = transforms.Compose([transforms.ToPILImage(), transforms.ToTensor()])
train_data = ImageData(df = labels, data_dir = train_dir, transform = data_transf)
train_loader = DataLoader(dataset = train_data, batch_size = 64)
from efficientnet_pytorch import EfficientNet
model = EfficientNet.from_name('efficientnet-b1')
for param in model.parameters():
param.requires_grad = True
num_ftrs = model._fc.in_features
model._fc = nn.Linear(num_ftrs, 1)
optimizer = optim.Adam(model.parameters())
for ii, (data, target) in enumerate(train_loader):
data, target = data.cuda(), target.cuda()
loss = loss_func(m(output), target)
loss_log.append(loss.item())
print('Epoch: {} - Loss: {:.6f}'.format(epoch + 1, loss.item()))